From e444e39fd08dde083c7cd9983002b40aea308e5d Mon Sep 17 00:00:00 2001
From: Pablo Lara <larabjj@gmail.com>
Date: Wed, 6 Nov 2024 08:29:39 +0100
Subject: [PATCH] feat: add helper function to monitor task state during
 execution

---
 .../table/data-table-row-actions.tsx          | 10 ++++-
 .../workflow/forms/test-connection-form.tsx   | 42 ++++++++++++-------
 lib/helper.ts                                 | 38 +++++++++++++++++
 3 files changed, 74 insertions(+), 16 deletions(-)

diff --git a/components/providers/table/data-table-row-actions.tsx b/components/providers/table/data-table-row-actions.tsx
index 127e829c4a..3e541751da 100644
--- a/components/providers/table/data-table-row-actions.tsx
+++ b/components/providers/table/data-table-row-actions.tsx
@@ -7,6 +7,7 @@ import {
   DropdownMenu,
   DropdownSection,
   DropdownTrigger,
+  Link,
 } from "@nextui-org/react";
 import {
   AddNoteBulkIcon,
@@ -35,6 +36,7 @@ export function DataTableRowActions<ProviderProps>({
   const [isEditOpen, setIsEditOpen] = useState(false);
   const [isDeleteOpen, setIsDeleteOpen] = useState(false);
   const providerId = (row.original as { id: string }).id;
+  const providerType = (row.original as { type: string }).type;
   const providerAlias = (row.original as any).attributes?.alias;
   return (
     <>
@@ -79,7 +81,13 @@ export function DataTableRowActions<ProviderProps>({
                 textValue="Check Connection"
                 startContent={<AddNoteBulkIcon className={iconClasses} />}
               >
-                {/* TODO: add the provider type to the search params */}
+                <Link
+                  href={`/providers/test-connection?type=${providerType}&id=${providerId}`}
+                >
+                  <span className="text-sm text-default-600">
+                    Test Connection
+                  </span>
+                </Link>
               </DropdownItem>
               <DropdownItem
                 key="edit"
diff --git a/components/providers/workflow/forms/test-connection-form.tsx b/components/providers/workflow/forms/test-connection-form.tsx
index 499b55ede8..3076ecb658 100644
--- a/components/providers/workflow/forms/test-connection-form.tsx
+++ b/components/providers/workflow/forms/test-connection-form.tsx
@@ -14,6 +14,7 @@ import { SaveIcon } from "@/components/icons";
 import { useToast } from "@/components/ui";
 import { CustomButton } from "@/components/ui/custom";
 import { Form } from "@/components/ui/form";
+import { checkTaskStatus } from "@/lib/helper";
 import { ApiError, testConnectionFormSchema } from "@/types";
 
 import { ProviderInfo } from "../..";
@@ -86,21 +87,30 @@ export const TestConnectionForm = ({
       const taskId = data.data.id;
       setApiErrorMessage(null);
 
-      const task = await getTask(taskId);
-      console.log({ task }, "task");
+      // Use the helper function to check the task status
+      const taskResult = await checkTaskStatus(taskId);
 
-      const connected = task.data.attributes.result.connected;
-      const error = task.data.attributes.result.error;
+      if (taskResult.completed) {
+        // If the task is completed, fetch the final task data
+        const task = await getTask(taskId);
+        const connected = task.data.attributes.result.connected;
 
-      setConnectionStatus({
-        connected,
-        error,
-      });
+        setConnectionStatus({
+          connected,
+          error: null,
+        });
 
-      if (connected) {
-        router.push(
-          `/providers/launch-scan?type=${providerType}&id=${providerId}`,
-        );
+        if (connected) {
+          router.push(
+            `/providers/launch-scan?type=${providerType}&id=${providerId}`,
+          );
+        }
+      } else {
+        // If the task failed, display the error message
+        setConnectionStatus({
+          connected: false,
+          error: taskResult.error || "Unknown error",
+        });
       }
     }
   };
@@ -115,9 +125,11 @@ export const TestConnectionForm = ({
           <div className="text-2xl font-bold leading-9 text-default-foreground">
             Test connection
           </div>
-          <div className="py-2 text-default-500">
-            Please check the provider connection
-          </div>
+          <p className="py-2 text-default-500">
+            Ensure all required credentials and configurations are completed
+            accurately. A successful connection will enable the option to
+            initiate a scan in the following step.
+          </p>
         </div>
 
         {apiErrorMessage && (
diff --git a/lib/helper.ts b/lib/helper.ts
index 2777e9d344..8e9eb843e6 100644
--- a/lib/helper.ts
+++ b/lib/helper.ts
@@ -1,7 +1,45 @@
+import { getTask } from "@/actions/task";
 import { MetaDataProps } from "@/types";
 
+export async function checkTaskStatus(
+  taskId: string,
+): Promise<{ completed: boolean; error?: string }> {
+  const MAX_RETRIES = 20; // Define the maximum number of attempts before stopping the polling
+  const RETRY_DELAY = 1000; // Delay time between each poll (in milliseconds)
+
+  for (let attempt = 0; attempt < MAX_RETRIES; attempt++) {
+    const task = await getTask(taskId);
+
+    if (task.error) {
+      console.error(`Error retrieving task: ${task.error}`);
+      return { completed: false, error: task.error };
+    }
+
+    const state = task.data.attributes.state;
+
+    switch (state) {
+      case "completed":
+        return { completed: true };
+      case "failed":
+        return { completed: false, error: task.data.attributes.result.error };
+      case "available":
+      case "scheduled":
+      case "executing":
+        // Continue waiting if the task is still in progress
+        await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY));
+        break;
+      default:
+        console.warn(`Unexpected task state: ${state}`);
+        return { completed: false, error: "Unexpected task state" };
+    }
+  }
+
+  return { completed: false, error: "Max retries exceeded" };
+}
+
 export const wait = (ms: number) =>
   new Promise((resolve) => setTimeout(resolve, ms));
+
 // Helper function to create dictionaries by type
 export function createDict(type: string, data: any) {
   const includedField = data?.included?.filter(