Update KindInput component to bring it inline with others

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

@@ -11,8 +11,7 @@ import { vscode } from "../vscode-api";
 
 import { Method } from "../../model-editor/method";
 import { ModeledMethod } from "../../model-editor/modeled-method";
-import { KindInput } from "./KindInput";
-import { extensiblePredicateDefinitions } from "../../model-editor/predicates";
+import { ModelKindDropdown } from "./ModelKindDropdown";
 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 (
@@ -154,12 +133,10 @@ function ModelableMethodRow(props: MethodRowProps) {
             />
           </VSCodeDataGridCell>
           <VSCodeDataGridCell gridColumn={5}>
-            <KindInput
-              kinds={predicate?.supportedKinds || []}
-              value={modeledMethod?.kind}
-              disabled={!showKindCell}
-              onChange={handleKindChange}
-              aria-label="Kind"
+            <ModelKindDropdown
+              method={method}
+              modeledMethod={modeledMethod}
+              onChange={onChange}
             />
           </VSCodeDataGridCell>
         </>

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

@@ -0,0 +1,84 @@
+import * as React from "react";
+import { ChangeEvent, useCallback, useEffect, useMemo } from "react";
+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 = {
+  method: Method;
+  modeledMethod: ModeledMethod | undefined;
+  onChange: (method: Method, modeledMethod: ModeledMethod) => void;
+};
+
+export const ModelKindDropdown = ({
+  method,
+  modeledMethod,
+  onChange,
+}: 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 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;
+
+      onChangeKind(kind);
+    },
+    [onChangeKind],
+  );
+
+  useEffect(() => {
+    const value = modeledMethod?.kind;
+    if (value === undefined && kinds.length > 0) {
+      onChangeKind(kinds[0]);
+    }
+
+    if (value !== undefined && !kinds.includes(value)) {
+      onChangeKind(kinds[0]);
+    }
+  }, [modeledMethod?.kind, kinds, onChangeKind]);
+
+  return (
+    <Dropdown
+      value={modeledMethod?.kind}
+      options={options}
+      disabled={disabled}
+      onChange={handleChange}
+      aria-label="Kind"
+    />
+  );
+};

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

@@ -0,0 +1,96 @@
+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
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("local");
+    await userEvent.selectOptions(screen.getByRole("combobox"), "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
+        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
+        method={method}
+        modeledMethod={updatedModeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    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
+        method={method}
+        modeledMethod={modeledMethod}
+        onChange={onChange}
+      />,
+    );
+
+    expect(screen.getByRole("combobox")).toHaveValue("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,
+  };
+}