From 8e231d7b1d20b160daef0ae02a2650b28689382a Mon Sep 17 00:00:00 2001
From: Koen Vlaswinkel <koesie10@github.com>
Date: Mon, 8 Apr 2024 14:48:53 +0200
Subject: [PATCH] Add unit tests for Python argument options

---
 .../model-editor/languages/python/index.ts    |  27 ++-
 .../languages/python/index.test.ts            | 208 ++++++++++++++++++
 2 files changed, 227 insertions(+), 8 deletions(-)
 create mode 100644 extensions/ql-vscode/test/unit-tests/model-editor/languages/python/index.test.ts

diff --git a/extensions/ql-vscode/src/model-editor/languages/python/index.ts b/extensions/ql-vscode/src/model-editor/languages/python/index.ts
index a9d57bb5d..946aac294 100644
--- a/extensions/ql-vscode/src/model-editor/languages/python/index.ts
+++ b/extensions/ql-vscode/src/model-editor/languages/python/index.ts
@@ -177,11 +177,28 @@ export const python: ModelsAsDataLanguage = {
     // Argument and Parameter are equivalent in Python, but we'll use Argument in the model editor
     const argumentsList = getArgumentsList(method.methodParameters).map(
       (argument, index): MethodArgument => {
+        if (
+          method.endpointType === EndpointType.Method &&
+          argument === "self" &&
+          index === 0
+        ) {
+          return {
+            path: "Argument[self]",
+            label: "Argument[self]: self",
+          };
+        }
+
+        // If this is a method, self does not count as an argument index, so we
+        // should start at 0 for the second argument
+        if (method.endpointType === EndpointType.Method) {
+          index -= 1;
+        }
+
         // Keyword-only arguments end with `:` in the query
         if (argument.endsWith(":")) {
           return {
             path: `Argument[${argument}]`,
-            label: `Argument[${argument}]`,
+            label: `Argument[${argument}]: ${argument.substring(0, argument.length - 1)}`,
           };
         }
 
@@ -202,13 +219,7 @@ export const python: ModelsAsDataLanguage = {
     );
 
     return {
-      options: [
-        {
-          path: "Argument[self]",
-          label: "Argument[self]",
-        },
-        ...argumentsList,
-      ],
+      options: argumentsList,
       // If there are no arguments, we will default to "Argument[self]"
       defaultArgumentPath:
         argumentsList.length > 0 ? argumentsList[0].path : "Argument[self]",
diff --git a/extensions/ql-vscode/test/unit-tests/model-editor/languages/python/index.test.ts b/extensions/ql-vscode/test/unit-tests/model-editor/languages/python/index.test.ts
new file mode 100644
index 000000000..ce2bfe0cf
--- /dev/null
+++ b/extensions/ql-vscode/test/unit-tests/model-editor/languages/python/index.test.ts
@@ -0,0 +1,208 @@
+import type { MethodDefinition } from "../../../../../src/model-editor/method";
+import { EndpointType } from "../../../../../src/model-editor/method";
+import { python } from "../../../../../src/model-editor/languages/python";
+import type { MethodArgumentOptions } from "../../../../../src/model-editor/languages";
+
+const testCases: Array<{
+  method: MethodDefinition;
+  options: MethodArgumentOptions;
+}> = [
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(a,b,c)",
+      endpointType: EndpointType.Function,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[0,a:]",
+          label: "Argument[0,a:]: a",
+        },
+        {
+          path: "Argument[1,b:]",
+          label: "Argument[1,b:]: b",
+        },
+        {
+          path: "Argument[2,c:]",
+          label: "Argument[2,c:]: c",
+        },
+      ],
+      defaultArgumentPath: "Argument[0,a:]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(self,a,b,c)",
+      endpointType: EndpointType.Method,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]: self",
+        },
+        {
+          path: "Argument[0,a:]",
+          label: "Argument[0,a:]: a",
+        },
+        {
+          path: "Argument[1,b:]",
+          label: "Argument[1,b:]: b",
+        },
+        {
+          path: "Argument[2,c:]",
+          label: "Argument[2,c:]: c",
+        },
+      ],
+      defaultArgumentPath: "Argument[self]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(a,b,c:)",
+      endpointType: EndpointType.Function,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[0,a:]",
+          label: "Argument[0,a:]: a",
+        },
+        {
+          path: "Argument[1,b:]",
+          label: "Argument[1,b:]: b",
+        },
+        {
+          path: "Argument[c:]",
+          label: "Argument[c:]: c",
+        },
+      ],
+      defaultArgumentPath: "Argument[0,a:]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(a/,b,c:)",
+      endpointType: EndpointType.Function,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[0]",
+          label: "Argument[0]: a",
+        },
+        {
+          path: "Argument[1,b:]",
+          label: "Argument[1,b:]: b",
+        },
+        {
+          path: "Argument[c:]",
+          label: "Argument[c:]: c",
+        },
+      ],
+      defaultArgumentPath: "Argument[0]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(self,a/,b/,c,d,e,f:,g:,h:)",
+      endpointType: EndpointType.Method,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]: self",
+        },
+        {
+          path: "Argument[0]",
+          label: "Argument[0]: a",
+        },
+        {
+          path: "Argument[1]",
+          label: "Argument[1]: b",
+        },
+        {
+          path: "Argument[2,c:]",
+          label: "Argument[2,c:]: c",
+        },
+        {
+          path: "Argument[3,d:]",
+          label: "Argument[3,d:]: d",
+        },
+        {
+          path: "Argument[4,e:]",
+          label: "Argument[4,e:]: e",
+        },
+        {
+          path: "Argument[f:]",
+          label: "Argument[f:]: f",
+        },
+        {
+          path: "Argument[g:]",
+          label: "Argument[g:]: g",
+        },
+        {
+          path: "Argument[h:]",
+          label: "Argument[h:]: h",
+        },
+      ],
+      defaultArgumentPath: "Argument[self]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "(self)",
+      endpointType: EndpointType.Method,
+    },
+    options: {
+      options: [
+        {
+          path: "Argument[self]",
+          label: "Argument[self]: self",
+        },
+      ],
+      defaultArgumentPath: "Argument[self]",
+    },
+  },
+  {
+    method: {
+      packageName: "requests",
+      typeName: "Session",
+      methodName: "foo",
+      methodParameters: "()",
+      endpointType: EndpointType.Function,
+    },
+    options: {
+      options: [],
+      defaultArgumentPath: "Argument[self]",
+    },
+  },
+];
+
+describe("getArgumentOptions", () => {
+  it.each(testCases)(
+    "returns the correct options for $method",
+    ({ method, options }) => {
+      expect(python.getArgumentOptions(method)).toEqual(options);
+    },
+  );
+});