From 754ab82a99009953b9ed25dd017a904ecf7d15b7 Mon Sep 17 00:00:00 2001
From: Arthur <worldnomad@gmail.com>
Date: Fri, 17 May 2024 11:16:06 -0700
Subject: [PATCH] 16145 allow both pk and script name

---
 netbox/extras/api/views.py      | 15 ++++++++++-----
 netbox/extras/tests/test_api.py | 16 ++++------------
 2 files changed, 14 insertions(+), 17 deletions(-)

diff --git a/netbox/extras/api/views.py b/netbox/extras/api/views.py
index 1be171dfa..89cf9c5df 100644
--- a/netbox/extras/api/views.py
+++ b/netbox/extras/api/views.py
@@ -217,12 +217,17 @@ class ScriptViewSet(ModelViewSet):
     lookup_value_regex = '[^/]+'  # Allow dots
 
     def _get_script(self, pk):
-        try:
-            module_name, script_name = pk.split('.', maxsplit=1)
-        except ValueError:
-            raise Http404
+        if pk.isnumeric():
+            script = get_object_or_404(self.queryset, pk=pk)
+            module = script.module
+        else:
+            try:
+                module_name, script_name = pk.split('.', maxsplit=1)
+            except ValueError:
+                raise Http404
+
+            module, script = get_module_and_script(module_name, script_name)
 
-        module, script = get_module_and_script(module_name, script_name)
         if script is None:
             raise Http404
 
diff --git a/netbox/extras/tests/test_api.py b/netbox/extras/tests/test_api.py
index 40635a936..5d243ae1a 100644
--- a/netbox/extras/tests/test_api.py
+++ b/netbox/extras/tests/test_api.py
@@ -777,23 +777,15 @@ class ScriptTest(APITestCase):
             is_executable=True,
         )
 
-    def get_vars(self, *args):
-        return {
-            k: v.__class__.__name__ for k, v in self.TestScriptClass._get_vars().items()
-        }
-
-    def get_test_script(self, *args):
-        module = ScriptModule.objects.first()
-        return module, module.scripts.first()
+    def python_class(self):
+        return self.TestScriptClass
 
     def setUp(self):
         super().setUp()
 
         # Monkey-patch the Script model to return our TestScriptClass above
         from extras.api.views import ScriptViewSet
-        ScriptViewSet._get_script = self.get_test_script
-        from extras.api.serializers_.scripts import ScriptSerializer
-        ScriptSerializer.get_vars = self.get_vars
+        Script.python_class = self.python_class
 
     def test_get_script(self):
         module = ScriptModule.objects.get(
@@ -801,7 +793,7 @@ class ScriptTest(APITestCase):
             file_path='/var/tmp/script.py'
         )
         script = module.scripts.all().first()
-        url = reverse('extras-api:script-detail', kwargs={'pk': None})
+        url = reverse('extras-api:script-detail', kwargs={'pk': script.pk})
         response = self.client.get(url, **self.header)
 
         self.assertEqual(response.data['name'], self.TestScriptClass.Meta.name)