Update ModelKindDropdown to be more self-contained

Author: charisk

extensions/ql-vscode/src/view/model-editor/MethodRow.tsx

@@ -12,7 +12,6 @@ import { vscode } from "../vscode-api";
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
 import { ModelKindDropdown } from "./ModelKindDropdown";
-import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
 import { Mode } from "../../model-editor/shared/mode";
 import { MethodClassifications } from "./MethodClassifications";
 import {
@@ -73,31 +72,11 @@ export const MethodRow = (props: MethodRowProps) => {
 function ModelableMethodRow(props: MethodRowProps) {
   const { method, modeledMethod, methodIsUnsaved, mode, onChange } = props;
 
-  const handleKindChange = useCallback(
-    (kind: string) => {
-      if (!modeledMethod) {
-        return;
-      }
-
-      onChange(method, {
-        ...modeledMethod,
-        kind,
-      });
-    },
-    [onChange, method, modeledMethod],
-  );
-
   const jumpToUsage = useCallback(
     () => sendJumpToUsageMessage(method),
     [method],
   );
 
-  const predicate =
-    modeledMethod?.type && modeledMethod.type !== "none"
-      ? extensiblePredicateDefinitions[modeledMethod.type]
-      : undefined;
-  const showKindCell = predicate?.supportedKinds;
-
   const modelingStatus = getModelingStatus(modeledMethod, methodIsUnsaved);
 
   return (
@@ -155,11 +134,9 @@ function ModelableMethodRow(props: MethodRowProps) {
           </VSCodeDataGridCell>
           <VSCodeDataGridCell gridColumn={5}>
             <ModelKindDropdown
-              kinds={predicate?.supportedKinds || []}
-              value={modeledMethod?.kind}
-              disabled={!showKindCell}
-              onChange={handleKindChange}
-              aria-label="Kind"
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
         </>

extensions/ql-vscode/src/view/model-editor/ModelKindDropdown.tsx

@@ -1,56 +1,84 @@
 import * as React from "react";
 import { ChangeEvent, useCallback, useEffect, useMemo } from "react";
-import type { ModeledMethodKind } from "../../model-editor/modeled-method";
+import type {
+  ModeledMethod,
+  ModeledMethodKind,
+} from "../../model-editor/modeled-method";
 import { Dropdown } from "../common/Dropdown";
+import { Method } from "../../model-editor/method";
+import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
 
 type Props = {
-  kinds: ModeledMethodKind[];
-
-  value: ModeledMethodKind | undefined;
-  disabled?: boolean;
-  onChange: (value: ModeledMethodKind) => void;
-
-  "aria-label"?: string;
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (method: Method, modeledMethod: ModeledMethod) => void;
 };
 
 export const ModelKindDropdown = ({
-  kinds,
-  value,
-  disabled,
+  method,
+  modeledMethod,
   onChange,
-  ...props
 }: Props) => {
+  const predicate = useMemo(() => {
+    return modeledMethod?.type && modeledMethod.type !== "none"
+      ? extensiblePredicateDefinitions[modeledMethod.type]
+      : undefined;
+  }, [modeledMethod?.type]);
+
+  const kinds = useMemo(() => predicate?.supportedKinds || [], [predicate]);
+
+  const disabled = useMemo(
+    () => !predicate?.supportedKinds,
+    [predicate?.supportedKinds],
+  );
+
   const options = useMemo(
     () => kinds.map((kind) => ({ value: kind, label: kind })),
     [kinds],
   );
 
-  const handleInput = useCallback(
+  const onChangeKind = useCallback(
+    (kind: ModeledMethodKind) => {
+      if (!modeledMethod) {
+        return;
+      }
+
+      onChange(method, {
+        ...modeledMethod,
+        kind,
+      });
+    },
+    [method, modeledMethod, onChange],
+  );
+
+  const handleChange = useCallback(
     (e: ChangeEvent<HTMLSelectElement>) => {
       const target = e.target as HTMLSelectElement;
+      const kind = target.value;
 
-      onChange(target.value);
+      onChangeKind(kind);
     },
-    [onChange],
+    [onChangeKind],
   );
 
   useEffect(() => {
+    const value = modeledMethod?.kind;
     if (value === undefined && kinds.length > 0) {
-      onChange(kinds[0]);
+      onChangeKind(kinds[0]);
     }
 
     if (value !== undefined && !kinds.includes(value)) {
-      onChange(kinds[0]);
+      onChangeKind(kinds[0]);
     }
-  }, [value, kinds, onChange]);
+  }, [modeledMethod?.kind, kinds, onChangeKind]);
 
   return (
     <Dropdown
-      value={value}
+      value={modeledMethod?.kind}
       options={options}
       disabled={disabled}
-      onChange={handleInput}
-      {...props}
+      onChange={handleChange}
+      aria-label="Kind"
     />
   );
 };

extensions/ql-vscode/src/view/model-editor/__tests__/ModelKindDropdown.spec.tsx

@@ -2,61 +2,95 @@ import * as React from "react";
 import { render, screen } from "@testing-library/react";
 import { ModelKindDropdown } from "../ModelKindDropdown";
 import userEvent from "@testing-library/user-event";
+import { createMethod } from "../../../../test/factories/model-editor/method-factories";
+import { createModeledMethod } from "../../../../test/factories/model-editor/modeled-method-factories";
 
 describe(ModelKindDropdown.name, () => {
   const onChange = jest.fn();
+  const method = createMethod();
 
   beforeEach(() => {
     onChange.mockReset();
   });
 
   it("allows changing the kind", async () => {
+    const modeledMethod = createModeledMethod({
+      type: "source",
+      kind: "local",
+    });
+
     render(
       <ModelKindDropdown
-        kinds={["local", "remote"]}
-        value="local"
+        method={method}
+        modeledMethod={modeledMethod}
         onChange={onChange}
       />,
     );
 
     expect(screen.getByRole("combobox")).toHaveValue("local");
     await userEvent.selectOptions(screen.getByRole("combobox"), "remote");
-    expect(onChange).toHaveBeenCalledWith("remote");
+    expect(onChange).toHaveBeenCalledWith(
+      method,
+      expect.objectContaining({
+        kind: "remote",
+      }),
+    );
   });
 
   it("resets the kind when changing the supported kinds", () => {
+    const method = createMethod();
+    const modeledMethod = createModeledMethod({
+      type: "source",
+      kind: "local",
+    });
+
     const { rerender } = render(
       <ModelKindDropdown
-        kinds={["local", "remote"]}
-        value={"local"}
+        method={method}
+        modeledMethod={modeledMethod}
         onChange={onChange}
       />,
     );
 
     expect(screen.getByRole("combobox")).toHaveValue("local");
     expect(onChange).not.toHaveBeenCalled();
 
+    // Changing the type to sink should update the supported kinds
+    const updatedModeledMethod = createModeledMethod({
+      type: "sink",
+    });
+
     rerender(
       <ModelKindDropdown
-        kinds={["sql-injection", "log-injection", "url-redirection"]}
-        value="local"
+        method={method}
+        modeledMethod={updatedModeledMethod}
         onChange={onChange}
       />,
     );
-    expect(screen.getByRole("combobox")).toHaveValue("sql-injection");
-    expect(onChange).toHaveBeenCalledWith("sql-injection");
+
+    expect(screen.getByRole("combobox")).toHaveValue("code-injection");
   });
 
   it("sets the kind when value is undefined", () => {
+    const method = createMethod();
+    const modeledMethod = createModeledMethod({
+      type: "source",
+    });
+
     render(
       <ModelKindDropdown
-        kinds={["local", "remote"]}
-        value={undefined}
+        method={method}
+        modeledMethod={modeledMethod}
         onChange={onChange}
       />,
     );
 
     expect(screen.getByRole("combobox")).toHaveValue("local");
-    expect(onChange).toHaveBeenCalledWith("local");
+    expect(onChange).toHaveBeenCalledWith(
+      method,
+      expect.objectContaining({
+        kind: "local",
+      }),
+    );
   });
 });

extensions/ql-vscode/test/factories/model-editor/modeled-method-factories.ts

@@ -0,0 +1,20 @@
+import { ModeledMethod } from "../../../src/model-editor/modeled-method";
+
+export function createModeledMethod(
+  data: Partial<ModeledMethod> = {},
+): ModeledMethod {
+  return {
+    libraryVersion: "1.6.0",
+    signature: "org.sql2o.Connection#createQuery(String)",
+    packageName: "org.sql2o",
+    typeName: "Connection",
+    methodName: "createQuery",
+    methodParameters: "(String)",
+    type: "sink",
+    input: "Argument[0]",
+    output: "",
+    kind: "jndi-injection",
+    provenance: "manual",
+    ...data,
+  };
+}