From ac69be2e8a8d0ac266ae820ac5280e7bb3410315 Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Thu, 28 Nov 2024 13:14:54 -0300
Subject: [PATCH] Fix: Disable the ComponentAsTool shortcut when the component
 code does not have tool_mode=True + regression tests (#4918)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

✨ (index.tsx): Add convertTestName function to generate data-testid attribute dynamically based on display_name
🔧 (use-shortcuts.tsx): Pass hasToolMode parameter to handleToolModeWShortcut function to conditionally activate tool mode
🔧 (index.tsx): Pass hasToolMode prop to NodeToolbarComponent to enable/disable tool mode functionality
🔧 (Vector Store.spec.ts): Import and use extractAndCleanCode function to extract and clean code content
🔧 (general-bugs-component-as-tool-shortcut.spec.ts): Import and use extractAndCleanCode function to extract and clean code content
📝 (extract-and-clean-code.ts): Add utility function extractAndCleanCode to extract and clean code content from HTML page
---
 .../sidebarDraggableComponent/index.tsx       |  4 +
 .../hooks/use-shortcuts.tsx                   |  9 +-
 .../components/nodeToolbarComponent/index.tsx |  1 +
 .../core/integrations/Vector Store.spec.ts    | 22 +----
 ...al-bugs-component-as-tool-shortcut.spec.ts | 88 +++++++++++++++++++
 .../tests/utils/extract-and-clean-code.ts     | 22 +++++
 6 files changed, 123 insertions(+), 23 deletions(-)
 create mode 100644 src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
 create mode 100644 src/frontend/tests/utils/extract-and-clean-code.ts

diff --git a/src/frontend/src/pages/FlowPage/components/flowSidebarComponent/components/sidebarDraggableComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/flowSidebarComponent/components/sidebarDraggableComponent/index.tsx
index ab3c6bb73..cdce9a7c3 100644
--- a/src/frontend/src/pages/FlowPage/components/flowSidebarComponent/components/sidebarDraggableComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/flowSidebarComponent/components/sidebarDraggableComponent/index.tsx
@@ -1,3 +1,4 @@
+import { convertTestName } from "@/components/common/storeCardComponent/utils/convert-test-name";
 import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
 import useDeleteFlow from "@/hooks/flows/use-delete-flow";
@@ -176,6 +177,9 @@ export const SidebarDraggableComponent = forwardRef(
               <div className="flex shrink-0 items-center gap-1">
                 {!disabled && (
                   <Button
+                    data-testid={`add-component-button-${convertTestName(
+                      display_name,
+                    )}`}
                     variant="ghost"
                     size="icon"
                     tabIndex={-1}
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/hooks/use-shortcuts.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/hooks/use-shortcuts.tsx
index 495bc33e3..0d4ad5cc2 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/hooks/use-shortcuts.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/hooks/use-shortcuts.tsx
@@ -18,6 +18,7 @@ export default function useShortcuts({
   ungroup,
   minimizeFunction,
   activateToolMode,
+  hasToolMode,
 }: {
   showOverrideModal?: boolean;
   showModalAdvanced?: boolean;
@@ -34,6 +35,7 @@ export default function useShortcuts({
   ungroup?: () => void;
   minimizeFunction?: () => void;
   activateToolMode?: () => void;
+  hasToolMode?: boolean;
 }) {
   const advancedSettings = useShortcutsStore((state) => state.advancedSettings);
   const minimize = useShortcutsStore((state) => state.minimize);
@@ -117,7 +119,8 @@ export default function useShortcuts({
     minimizeFunction();
   }
 
-  function handleToolModeWShortcut(e: KeyboardEvent) {
+  function handleToolModeWShortcut(e: KeyboardEvent, hasToolMode?: boolean) {
+    if (!hasToolMode) return;
     if (isWrappedWithClass(e, "noflow") || !activateToolMode) return;
     e.preventDefault();
     activateToolMode();
@@ -135,5 +138,7 @@ export default function useShortcuts({
   useHotkeys(download, handleDownloadWShortcut, { preventDefault: true });
   useHotkeys(freeze, handleFreeze);
   useHotkeys(freezeAll, handleFreezeAll);
-  useHotkeys(toolMode, handleToolModeWShortcut, { preventDefault: true });
+  useHotkeys(toolMode, (e) => handleToolModeWShortcut(e, hasToolMode), {
+    preventDefault: true,
+  });
 }
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 3c3170e32..3cd3f3303 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -189,6 +189,7 @@ export default function NodeToolbarComponent({
     ungroup: handleungroup,
     minimizeFunction: minimize,
     activateToolMode: activateToolMode,
+    hasToolMode,
   });
 
   const paste = useFlowStore((state) => state.paste);
diff --git a/src/frontend/tests/core/integrations/Vector Store.spec.ts b/src/frontend/tests/core/integrations/Vector Store.spec.ts
index d425e2807..9bd131cc8 100644
--- a/src/frontend/tests/core/integrations/Vector Store.spec.ts	
+++ b/src/frontend/tests/core/integrations/Vector Store.spec.ts	
@@ -1,6 +1,7 @@
 import { Page, test } from "@playwright/test";
 import path from "path";
 import uaParser from "ua-parser-js";
+import { extractAndCleanCode } from "../../utils/extract-and-clean-code";
 
 test(
   "Vector Store RAG",
@@ -192,24 +193,3 @@ test(
     await page.getByTestId("input-chat-playground").last().isVisible();
   },
 );
-
-async function extractAndCleanCode(page: Page): Promise<string> {
-  const outerHTML = await page
-    .locator('//*[@id="codeValue"]')
-    .evaluate((el) => el.outerHTML);
-
-  const valueMatch = outerHTML.match(/value="([\s\S]*?)"/);
-  if (!valueMatch) {
-    throw new Error("Could not find value attribute in the HTML");
-  }
-
-  let codeContent = valueMatch[1]
-    .replace(/&quot;/g, '"')
-    .replace(/&amp;/g, "&")
-    .replace(/&lt;/g, "<")
-    .replace(/&gt;/g, ">")
-    .replace(/&#x27;/g, "'")
-    .replace(/&#x2F;/g, "/");
-
-  return codeContent;
-}
diff --git a/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts b/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
new file mode 100644
index 000000000..a89a3b920
--- /dev/null
+++ b/src/frontend/tests/extended/regression/general-bugs-component-as-tool-shortcut.spec.ts
@@ -0,0 +1,88 @@
+import { expect, test } from "@playwright/test";
+import { extractAndCleanCode } from "../../utils/extract-and-clean-code";
+
+test(
+  "user must be able to use component as tool shortcut only if has tool mode is True",
+  { tag: ["@release", "@components"] },
+  async ({ page }) => {
+    await page.goto("/");
+    await page.waitForTimeout(1000);
+    let modalCount = 0;
+    try {
+      const modalTitleElement = await page?.getByTestId("modal-title");
+      if (modalTitleElement) {
+        modalCount = await modalTitleElement.count();
+      }
+    } catch (error) {
+      modalCount = 0;
+    }
+
+    while (modalCount === 0) {
+      await page.getByText("New Flow", { exact: true }).click();
+      await page.waitForSelector('[data-testid="modal-title"]', {
+        timeout: 3000,
+      });
+      modalCount = await page.getByTestId("modal-title")?.count();
+    }
+
+    await page.getByTestId("blank-flow").click();
+
+    await page.waitForSelector('[data-testid="fit_view"]', {
+      timeout: 100000,
+    });
+    await page.getByTestId("sidebar-search-input").click();
+    await page.getByTestId("sidebar-search-input").fill("id generator");
+
+    await page
+      .getByTestId("helpersID Generator")
+      .hover()
+      .then(async () => {
+        await page.getByTestId("add-component-button-id-generator").click();
+      });
+
+    await page.waitForSelector('[data-testid="title-ID Generator"]', {
+      timeout: 3000,
+    });
+
+    expect(await page.getByText("Toolset", { exact: true }).count()).toBe(0);
+
+    await page.getByTestId("title-ID Generator").click();
+    await page.keyboard.press("ControlOrMeta+Shift+m");
+
+    expect(await page.getByText("Toolset", { exact: true }).count()).toBe(0);
+
+    await page.getByTestId("title-ID Generator").click();
+
+    await page.waitForSelector('[data-testid="code-button-modal"]', {
+      timeout: 3000,
+    });
+
+    await page.getByTestId("code-button-modal").click();
+
+    let code = await extractAndCleanCode(page);
+    code = code!.replace(
+      "refresh_button=True,",
+      "refresh_button=True,\n tool_mode=True,",
+    );
+
+    await page.locator("textarea").last().press(`ControlOrMeta+a`);
+    await page.keyboard.press("Backspace");
+    await page.locator("textarea").last().fill(code);
+    await page.locator('//*[@id="checkAndSaveBtn"]').click();
+
+    await page.waitForSelector('[data-testid="fit_view"]', {
+      timeout: 3000,
+    });
+
+    await page.getByTestId("title-ID Generator").click();
+    await page.keyboard.press("ControlOrMeta+Shift+m");
+
+    await page.waitForSelector('text="Toolset"', {
+      timeout: 3000,
+    });
+
+    expect(
+      await page.getByText("Toolset", { exact: true }).count(),
+    ).toBeGreaterThan(0);
+  },
+);
diff --git a/src/frontend/tests/utils/extract-and-clean-code.ts b/src/frontend/tests/utils/extract-and-clean-code.ts
new file mode 100644
index 000000000..95ca33da7
--- /dev/null
+++ b/src/frontend/tests/utils/extract-and-clean-code.ts
@@ -0,0 +1,22 @@
+import { Page } from "playwright/test";
+
+export async function extractAndCleanCode(page: Page): Promise<string> {
+  const outerHTML = await page
+    .locator('//*[@id="codeValue"]')
+    .evaluate((el) => el.outerHTML);
+
+  const valueMatch = outerHTML.match(/value="([\s\S]*?)"/);
+  if (!valueMatch) {
+    throw new Error("Could not find value attribute in the HTML");
+  }
+
+  let codeContent = valueMatch[1]
+    .replace(/&quot;/g, '"')
+    .replace(/&amp;/g, "&")
+    .replace(/&lt;/g, "<")
+    .replace(/&gt;/g, ">")
+    .replace(/&#x27;/g, "'")
+    .replace(/&#x2F;/g, "/");
+
+  return codeContent;
+}