From 27e3b6f377afb8c88e2e79ca805c335d15cff11a Mon Sep 17 00:00:00 2001
From: dansheps <dans@dansheps.com>
Date: Thu, 27 Feb 2020 07:45:11 -0600
Subject: [PATCH] Remove second variables, make widget mandatory on ObjectVar
 and MultiObjectVar

---
 netbox/extras/scripts.py            | 28 +++-----------
 netbox/extras/tests/test_scripts.py | 57 +++--------------------------
 2 files changed, 11 insertions(+), 74 deletions(-)

diff --git a/netbox/extras/scripts.py b/netbox/extras/scripts.py
index eb796cf32..fdde58a83 100644
--- a/netbox/extras/scripts.py
+++ b/netbox/extras/scripts.py
@@ -26,8 +26,6 @@ __all__ = [
     'BaseScript',
     'BooleanVar',
     'ChoiceVar',
-    'DynamicObjectVar',
-    'DynamicMultiObjectVar',
     'FileVar',
     'IntegerVar',
     'IPAddressVar',
@@ -170,10 +168,10 @@ class ObjectVar(ScriptVariable):
     """
     NetBox object representation. The provided QuerySet will determine the choices available.
     """
-    form_field = forms.ModelChoiceField
+    form_field = DynamicModelChoiceField
 
-    def __init__(self, queryset, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, queryset, widget, *args, **kwargs):
+        super().__init__(widget=widget, *args, **kwargs)
 
         # Queryset for field choices
         self.field_attrs['queryset'] = queryset
@@ -187,10 +185,10 @@ class MultiObjectVar(ScriptVariable):
     """
     Like ObjectVar, but can represent one or more objects.
     """
-    form_field = forms.ModelMultipleChoiceField
+    form_field = DynamicModelMultipleChoiceField
 
-    def __init__(self, queryset, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self, queryset, widget, *args, **kwargs):
+        super().__init__(widget=widget, *args, **kwargs)
 
         # Queryset for field choices
         self.field_attrs['queryset'] = queryset
@@ -200,20 +198,6 @@ class MultiObjectVar(ScriptVariable):
             self.form_field = TreeNodeMultipleChoiceField
 
 
-class DynamicObjectVar(ObjectVar):
-    """
-    A dynamic netbox object variable.  APISelect will determine the available choices
-    """
-    form_field = DynamicModelChoiceField
-
-
-class DynamicMultiObjectVar(MultiObjectVar):
-    """
-    A multiple choice version of DynamicObjectVar
-    """
-    form_field = DynamicModelMultipleChoiceField
-
-
 class FileVar(ScriptVariable):
     """
     An uploaded file.
diff --git a/netbox/extras/tests/test_scripts.py b/netbox/extras/tests/test_scripts.py
index c926a0b29..2d4d5b3fd 100644
--- a/netbox/extras/tests/test_scripts.py
+++ b/netbox/extras/tests/test_scripts.py
@@ -4,6 +4,7 @@ from netaddr import IPAddress, IPNetwork
 
 from dcim.models import DeviceRole
 from extras.scripts import *
+from utilities.forms import APISelect, APISelectMultiple
 
 
 class ScriptVariablesTest(TestCase):
@@ -129,31 +130,8 @@ class ScriptVariablesTest(TestCase):
         class TestScript(Script):
 
             var1 = ObjectVar(
-                queryset=DeviceRole.objects.all()
-            )
-
-        # Populate some objects
-        for i in range(1, 6):
-            DeviceRole(
-                name='Device Role {}'.format(i),
-                slug='device-role-{}'.format(i)
-            ).save()
-
-        # Validate valid data
-        data = {'var1': DeviceRole.objects.first().pk}
-        form = TestScript().as_form(data, None)
-        self.assertTrue(form.is_valid())
-        self.assertEqual(form.cleaned_data['var1'].pk, data['var1'])
-
-    def test_dynamicobjectvar(self):
-        """
-        Test dynamic version of the objectvar
-        """
-
-        class TestScript(Script):
-
-            var1 = DynamicObjectVar(
-                queryset=DeviceRole.objects.all()
+                queryset=DeviceRole.objects.all(),
+                widget=APISelect(api_url='/api/dcim/device-roles/')
             )
 
         # Populate some objects
@@ -174,33 +152,8 @@ class ScriptVariablesTest(TestCase):
         class TestScript(Script):
 
             var1 = MultiObjectVar(
-                queryset=DeviceRole.objects.all()
-            )
-
-        # Populate some objects
-        for i in range(1, 6):
-            DeviceRole(
-                name='Device Role {}'.format(i),
-                slug='device-role-{}'.format(i)
-            ).save()
-
-        # Validate valid data
-        data = {'var1': [role.pk for role in DeviceRole.objects.all()[:3]]}
-        form = TestScript().as_form(data, None)
-        self.assertTrue(form.is_valid())
-        self.assertEqual(form.cleaned_data['var1'][0].pk, data['var1'][0])
-        self.assertEqual(form.cleaned_data['var1'][1].pk, data['var1'][1])
-        self.assertEqual(form.cleaned_data['var1'][2].pk, data['var1'][2])
-
-    def test_dynamicmultiobjectvar(self):
-        """
-        Test dynamic version of the multiobjectvar
-        """
-
-        class TestScript(Script):
-
-            var1 = DynamicMultiObjectVar(
-                queryset=DeviceRole.objects.all()
+                queryset=DeviceRole.objects.all(),
+                widget=APISelectMultiple(api_url='/api/dcim/device-roles/')
             )
 
         # Populate some objects