Add module-aware template lookup utility

Author: SaravanakumarR2018

src/frontend/src/CustomNodes/helpers/check-code-validity.ts

@@ -1,5 +1,6 @@
 import { componentsToIgnoreUpdate } from "@/constants/constants";
 import type { OutputFieldType } from "@/types/api";
+import { resolveTemplateEntry } from "@/utils/component-template-utils";
 import type { NodeDataType } from "../../types/flow";
 
 // Returns true if the code is outdated (code string changed and not ignored)
@@ -53,12 +54,19 @@ const codeHasBreakingChange = (
 export const checkCodeValidity = (
   data: NodeDataType,
   templates: { [key: string]: any },
+  templatesByModule?: { [key: string]: any },
 ) => {
   if (!data?.node || !templates) return;
-  const template = templates[data.type]?.template;
+  const templateEntry = resolveTemplateEntry(
+    data.node,
+    templates,
+    templatesByModule,
+    data.type,
+  );
+  const template = templateEntry?.template;
   const currentCode = template?.code?.value;
   const thisNodesCode = data.node!.template?.code?.value;
-  const originalOutputs = templates[data.type]?.outputs;
+  const originalOutputs = templateEntry?.outputs;
   const userOutputs = data.node?.outputs;
   const originalTemplate = template;
   const userTemplate = data.node?.template;

src/frontend/src/stores/__tests__/flowStore.test.ts

@@ -83,6 +83,7 @@ jest.mock("../typesStore", () => ({
   useTypesStore: {
     getState: () => ({
       templates: {},
+      templatesByModule: {},
       types: {},
     }),
   },

src/frontend/src/stores/__tests__/typesStore.test.ts

@@ -23,11 +23,17 @@ jest.mock("../../utils/reactflowUtils", () => ({
   }),
 }));
 
+jest.mock("../../utils/component-template-utils", () => ({
+  buildTemplatesByModule: jest.fn(() => ({})),
+}));
+
 // Mock imports
 const mockExtractSecretFieldsFromComponents =
   require("../../utils/reactflowUtils").extractSecretFieldsFromComponents;
 const mockTemplatesGenerator =
   require("../../utils/reactflowUtils").templatesGenerator;
+const mockBuildTemplatesByModule =
+  require("../../utils/component-template-utils").buildTemplatesByModule;
 const mockTypesGenerator = require("../../utils/reactflowUtils").typesGenerator;
 
 const mockAPIData: APIDataType = {
@@ -85,13 +91,15 @@ describe("useTypesStore", () => {
       return new Set(Object.keys(data));
     });
     mockTemplatesGenerator.mockReturnValue({});
+    mockBuildTemplatesByModule.mockReturnValue({});
     mockTypesGenerator.mockReturnValue({});
 
     act(() => {
       useTypesStore.setState({
         ComponentFields: new Set(),
         types: {},
         templates: {},
+        templatesByModule: {},
         data: {},
       });
     });
@@ -104,6 +112,7 @@ describe("useTypesStore", () => {
       expect(result.current.ComponentFields).toEqual(new Set());
       expect(result.current.types).toEqual({});
       expect(result.current.templates).toEqual({});
+      expect(result.current.templatesByModule).toEqual({});
       expect(result.current.data).toEqual({});
     });
   });
@@ -209,6 +218,9 @@ describe("useTypesStore", () => {
     it("should set types, templates, data, and component fields", () => {
       mockTypesGenerator.mockReturnValue(mockTypes);
       mockTemplatesGenerator.mockReturnValue(mockTemplates);
+      mockBuildTemplatesByModule.mockReturnValue({
+        "module.text": mockTemplates.TextInput,
+      });
       mockExtractSecretFieldsFromComponents.mockReturnValue(
         new Set(["TextInput", "NumberInput"]),
       );
@@ -221,12 +233,16 @@ describe("useTypesStore", () => {
 
       expect(mockTypesGenerator).toHaveBeenCalledWith(mockAPIData);
       expect(mockTemplatesGenerator).toHaveBeenCalledWith(mockAPIData);
+      expect(mockBuildTemplatesByModule).toHaveBeenCalledWith(mockAPIData);
       expect(mockExtractSecretFieldsFromComponents).toHaveBeenCalledWith(
         mockAPIData,
       );
 
       expect(result.current.types).toEqual(mockTypes);
       expect(result.current.templates).toEqual(mockTemplates);
+      expect(result.current.templatesByModule).toEqual({
+        "module.text": mockTemplates.TextInput,
+      });
       expect(result.current.data).toEqual(mockAPIData);
       expect(result.current.ComponentFields).toEqual(
         new Set(["TextInput", "NumberInput"]),

src/frontend/src/stores/flowStore.ts

@@ -91,9 +91,14 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
   updateComponentsToUpdate: (nodes) => {
     const outdatedNodes: ComponentsToUpdateType[] = [];
     const templates = useTypesStore.getState().templates;
+    const templatesByModule = useTypesStore.getState().templatesByModule;
     nodes.forEach((node) => {
       if (node.type === "genericNode") {
-        const codeValidity = checkCodeValidity(node.data, templates);
+        const codeValidity = checkCodeValidity(
+          node.data,
+          templates,
+          templatesByModule,
+        );
         if (codeValidity && codeValidity.outdated)
           outdatedNodes.push({
             id: node.id,

src/frontend/src/stores/typesStore.ts

@@ -6,6 +6,7 @@ import {
   templatesGenerator,
   typesGenerator,
 } from "../utils/reactflowUtils";
+import { buildTemplatesByModule } from "../utils/component-template-utils";
 
 export const useTypesStore = create<TypesStoreType>((set, get) => ({
   ComponentFields: new Set(),
@@ -17,6 +18,7 @@ export const useTypesStore = create<TypesStoreType>((set, get) => ({
   },
   types: {},
   templates: {},
+  templatesByModule: {},
   data: {},
   setTypes: (data: APIDataType) => {
     set((old) => ({
@@ -27,6 +29,7 @@ export const useTypesStore = create<TypesStoreType>((set, get) => ({
         ...data,
       }),
       templates: templatesGenerator(data),
+      templatesByModule: buildTemplatesByModule(data),
     }));
   },
   setTemplates: (newState: {}) => {

src/frontend/src/types/zustand/types/index.ts

@@ -4,6 +4,7 @@ export type TypesStoreType = {
   types: { [char: string]: string };
   setTypes: (newState: {}) => void;
   templates: { [char: string]: APIClassType };
+  templatesByModule: { [module: string]: APIClassType };
   setTemplates: (newState: {}) => void;
   data: APIDataType;
   setData: (newState: {}) => void;

src/frontend/src/utils/component-template-utils.ts

@@ -0,0 +1,31 @@
+import type { APIClassType, APIKindType, APIObjectType } from "@/types/api";
+
+export function buildTemplatesByModule(
+  data: APIObjectType,
+): Record<string, APIClassType> {
+  return Object.keys(data).reduce((acc, curr) => {
+    Object.values(data[curr]).forEach((component: APIKindType[keyof APIKindType]) => {
+      const moduleName = component?.metadata?.module;
+      if (moduleName && !acc[moduleName]) {
+        acc[moduleName] = component as APIClassType;
+      }
+    });
+    return acc;
+  }, {} as Record<string, APIClassType>);
+}
+
+export function resolveTemplateEntry(
+  node: { metadata?: { module?: string } } | undefined,
+  templatesByName: Record<string, APIClassType>,
+  templatesByModule?: Record<string, APIClassType>,
+  typeKey?: string,
+): APIClassType | undefined {
+  const moduleName = node?.metadata?.module;
+  if (moduleName && templatesByModule?.[moduleName]) {
+    return templatesByModule[moduleName];
+  }
+  if (typeKey) {
+    return templatesByName[typeKey];
+  }
+  return undefined;
+}

src/frontend/src/utils/reactflowUtils.ts

@@ -1838,6 +1838,18 @@ export function templatesGenerator(data: APIObjectType) {
   }, {});
 }
 
+export function templatesByModuleGenerator(data: APIObjectType) {
+  return Object.keys(data).reduce((acc, curr) => {
+    Object.values(data[curr]).forEach((component) => {
+      const moduleName = component?.metadata?.module;
+      if (moduleName && !acc[moduleName]) {
+        acc[moduleName] = component;
+      }
+    });
+    return acc;
+  }, {} as Record<string, APIClassType>);
+}
+
 /**
  * Determines if a field is a SecretStr field type
  */