From 7641e81ca3b38db2df38d2617e0abcbc63ebde9b Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Wed, 28 Aug 2024 16:02:42 -0300
Subject: [PATCH] fix: add retriever input_type on RetrieverToolComponent +
 tests (#3601)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* ✨ (RetrieverTool.py): add "input_types" field to the configuration of the RetrieverToolComponent class

* ✅ (generalBugs-shard-12.spec.ts): add end-to-end test for connecting RetrieverTool to other components in the application.
---
 .../components/tools/RetrieverTool.py         |  1 +
 .../generalBugs-shard-12.spec.ts              | 97 +++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 src/frontend/tests/scheduled-end-to-end/generalBugs-shard-12.spec.ts

diff --git a/src/backend/base/langflow/components/tools/RetrieverTool.py b/src/backend/base/langflow/components/tools/RetrieverTool.py
index 43dfbcba636f..e3a7c2a7cfd8 100644
--- a/src/backend/base/langflow/components/tools/RetrieverTool.py
+++ b/src/backend/base/langflow/components/tools/RetrieverTool.py
@@ -15,6 +15,7 @@ def build_config(self):
                 "display_name": "Retriever",
                 "info": "Retriever to interact with",
                 "type": BaseRetriever,
+                "input_types": ["Retriever"],
             },
             "name": {"display_name": "Name", "info": "Name of the tool"},
             "description": {"display_name": "Description", "info": "Description of the tool"},
diff --git a/src/frontend/tests/scheduled-end-to-end/generalBugs-shard-12.spec.ts b/src/frontend/tests/scheduled-end-to-end/generalBugs-shard-12.spec.ts
new file mode 100644
index 000000000000..f1ab131a2888
--- /dev/null
+++ b/src/frontend/tests/scheduled-end-to-end/generalBugs-shard-12.spec.ts
@@ -0,0 +1,97 @@
+import { expect, test } from "@playwright/test";
+
+test("user should be able to connect RetrieverTool to another components", async ({
+  page,
+}) => {
+  await page.goto("/");
+  await page.waitForSelector('[data-testid="mainpage_title"]', {
+    timeout: 30000,
+  });
+
+  await page.waitForSelector('[id="new-project-btn"]', {
+    timeout: 30000,
+  });
+
+  let modalCount = 0;
+  try {
+    const modalTitleElement = await page.getByTestId("modal-title");
+    modalCount = await modalTitleElement.count();
+  } catch (error) {
+    modalCount = 0;
+  }
+
+  while (modalCount === 0) {
+    await page.getByText("New Project", { exact: true }).click();
+    await page.waitForTimeout(3000);
+    modalCount = await page.getByTestId("modal-title").count();
+  }
+
+  await page.waitForSelector('[data-testid="blank-flow"]', {
+    timeout: 30000,
+  });
+  await page.getByTestId("blank-flow").click();
+  await page.waitForSelector('[data-testid="extended-disclosure"]', {
+    timeout: 30000,
+  });
+  await page.getByTestId("extended-disclosure").click();
+  await page.getByPlaceholder("Search").click();
+  await page.getByPlaceholder("Search").fill("retriever");
+
+  await page.waitForTimeout(1000);
+
+  let modelElement = await page.getByTestId("toolsRetrieverTool");
+  let targetElement = await page.locator('//*[@id="react-flow-id"]');
+  await modelElement.dragTo(targetElement);
+
+  await page.mouse.up();
+  await page.mouse.down();
+
+  await page.getByTitle("fit view").click();
+  await page.getByTitle("zoom out").click();
+
+  await page.waitForTimeout(1000);
+
+  await page.getByTitle("zoom out").click();
+  await page
+    .locator('//*[@id="react-flow-id"]')
+    .hover()
+    .then(async () => {
+      await page.mouse.down();
+      await page.mouse.move(-500, 500);
+    });
+
+  await page.mouse.up();
+
+  await page.getByPlaceholder("Search").click();
+  await page.getByPlaceholder("Search").fill("chroma");
+
+  await page.waitForTimeout(1000);
+
+  modelElement = await page.getByTestId("vectorstoresChroma DB");
+  targetElement = await page.locator('//*[@id="react-flow-id"]');
+  await modelElement.dragTo(targetElement);
+
+  await page.mouse.up();
+  await page.mouse.down();
+
+  await page.waitForTimeout(1000);
+
+  await page.getByTitle("fit view").click();
+  await page.getByTitle("fit view").click();
+
+  //connection
+  const chromaDbOutput = await page
+    .getByTestId("handle-chroma-shownode-retriever-right")
+    .nth(0);
+  await chromaDbOutput.hover();
+  await page.mouse.down();
+  const retrieverToolInput = await page
+    .getByTestId("handle-retrievertool-shownode-retriever-left")
+    .nth(0);
+  await retrieverToolInput.hover();
+  await page.mouse.up();
+
+  await page.waitForTimeout(1000);
+
+  expect(await page.locator(".react-flow__edge-interaction").count()).toBe(1);
+});