Add tests for multi-value ScriptVariables

netbox/extras/tests/test_scripts.py

@@ -1,5 +1,7 @@
 from django.core.files.uploadedfile import SimpleUploadedFile
+from django.http import QueryDict
 from django.test import TestCase
+from django.utils.datastructures import MultiValueDict
 from netaddr import IPAddress, IPNetwork
 
 from dcim.models import DeviceRole
@@ -145,6 +147,33 @@ class ScriptVariablesTest(TestCase):
         self.assertTrue(form.is_valid())
         self.assertEqual(form.cleaned_data['var1'].pk, data['var1'])
 
+    def test_objectvar_initial(self):
+
+        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 initial data
+        keys = list(DeviceRole.objects.values_list('pk', flat=True)[:2])
+
+        initial = QueryDict(f"var1={keys[0]}&var1={keys[1]}")
+        form = TestScript().as_form(initial=initial)
+
+        # For a single-value field the initial value must be the LAST value in the query-string
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1'),
+            [keys[1], str(keys[1])]  # Accept both numeric and string values
+        )
+
     def test_multiobjectvar(self):
 
         class TestScript(Script):
@@ -168,6 +197,37 @@ class ScriptVariablesTest(TestCase):
         self.assertEqual(form.cleaned_data['var1'][1].pk, data['var1'][1])
         self.assertEqual(form.cleaned_data['var1'][2].pk, data['var1'][2])
 
+    def test_multiobjectvar_initial(self):
+
+        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 initial data
+        keys = list(DeviceRole.objects.values_list('pk', flat=True)[:2])
+
+        initial = QueryDict(f"var1={keys[0]}&var1={keys[1]}")
+        form = TestScript().as_form(initial=initial)
+
+        self.assertEqual(len(form.get_initial_for_field(form.fields['var1'], 'var1')), 2)
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1')[0],
+            [keys[0], str(keys[0])]  # Accept both numeric and string values
+        )
+        self.assertIn(
+            form.get_initial_for_field(form.fields['var1'], 'var1')[1],
+            [keys[1], str(keys[1])]  # Accept both numeric and string values
+        )
+
     def test_filevar(self):
 
         class TestScript(Script):