chore(ui): improve test connection workflow delete flow (#11912)

* chore(ui): improve test connection workflow delete flow

* delete workflow on component unmount

* add unit test

* chore: simplify the condition

* chore: refactor delete workflow method

* fix: alignment issue on server connection tab

---------

Co-authored-by: Chirag Madlani <12962843+chirag-madlani@users.noreply.github.com>

openmetadata-ui/src/main/resources/ui/src/components/common/AirflowMessageBanner/AirflowMessageBanner.tsx

@@ -18,9 +18,9 @@ import React, { FC } from 'react';
 import './airflow-message-banner.less';
 
 const AirflowMessageBanner: FC<SpaceProps> = ({ className }) => {
-  const { reason, isAirflowAvailable } = useAirflowStatus();
+  const { reason, isAirflowAvailable, isFetchingStatus } = useAirflowStatus();
 
-  if (isAirflowAvailable) {
+  if (isAirflowAvailable || isFetchingStatus) {
     return null;
   }
 

openmetadata-ui/src/main/resources/ui/src/components/common/TestConnection/TestConnection.test.tsx

@@ -27,6 +27,7 @@ import { ConfigData } from 'interface/service.interface';
 import React from 'react';
 import {
   addWorkflow,
+  deleteWorkflowById,
   getTestConnectionDefinitionByName,
   getWorkflowById,
   triggerWorkflowById,
@@ -476,4 +477,29 @@ describe('Test Connection Component', () => {
 
     expect(addWorkflow).not.toHaveBeenCalled();
   });
+
+  it('Should delete the workflow if workflow failed to trigger', async () => {
+    (triggerWorkflowById as jest.Mock).mockImplementationOnce(() =>
+      Promise.resolve(204)
+    );
+    jest.useFakeTimers();
+    await act(async () => {
+      render(<TestConnection {...mockProps} />);
+    });
+
+    const testConnectionButton = screen.getByTestId('test-connection-btn');
+
+    await act(async () => {
+      userEvent.click(testConnectionButton);
+    });
+
+    expect(addWorkflow).toHaveBeenCalledWith(CREATE_WORKFLOW_PAYLOAD);
+
+    expect(triggerWorkflowById).toHaveBeenCalledWith(WORKFLOW_DETAILS.id);
+
+    jest.advanceTimersByTime(2000);
+
+    // delete api should be called
+    expect(deleteWorkflowById).toHaveBeenCalledWith(WORKFLOW_DETAILS.id, true);
+  });
 });

openmetadata-ui/src/main/resources/ui/src/components/common/TestConnection/TestConnection.tsx

@@ -23,7 +23,7 @@ import {
   WorkflowType,
 } from 'generated/entity/automations/workflow';
 import { TestConnectionStep } from 'generated/entity/services/connections/testConnectionDefinition';
-import { toNumber } from 'lodash';
+import { isEmpty, toNumber } from 'lodash';
 import React, { FC, useEffect, useMemo, useRef, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import {
@@ -169,9 +169,12 @@ const TestConnection: FC<TestConnectionProps> = ({
   };
 
   const handleDeleteWorkflow = async (workflowId: string) => {
+    if (isEmpty(workflowId)) {
+      return;
+    }
+
     try {
-      const response = await deleteWorkflowById(workflowId, true);
+      await deleteWorkflowById(workflowId, true);
-      setCurrentWorkflow(response);
     } catch (error) {
       // do not throw error for this API
     }
@@ -218,6 +221,9 @@ const TestConnection: FC<TestConnectionProps> = ({
         setMessage(TEST_CONNECTION_FAILURE_MESSAGE);
         setIsTestingConnection(false);
 
+        // delete the workflow if workflow is not triggered successfully
+        await handleDeleteWorkflow(response.id);
+
         return;
       }
 
@@ -301,6 +307,12 @@ const TestConnection: FC<TestConnectionProps> = ({
       setMessage(TEST_CONNECTION_FAILURE_MESSAGE);
       setTestStatus(StatusType.Failed);
       showErrorToast(error as AxiosError);
+
+      // delete the workflow if there is an exception
+      const workflowId = currentWorkflowRef.current?.id;
+      if (workflowId) {
+        await handleDeleteWorkflow(workflowId);
+      }
     }
   };
 
@@ -320,6 +332,18 @@ const TestConnection: FC<TestConnectionProps> = ({
     currentWorkflowRef.current = currentWorkflow; // update ref with latest value of currentWorkflow state variable
   }, [currentWorkflow]);
 
+  useEffect(() => {
+    return () => {
+      /**
+       * if workflow is present then delete the workflow when component unmount
+       */
+      const workflowId = currentWorkflowRef.current?.id;
+      if (workflowId) {
+        handleDeleteWorkflow(workflowId);
+      }
+    };
+  }, []);
+
   // rendering
 
   return (

openmetadata-ui/src/main/resources/ui/src/pages/service/index.tsx

@@ -900,9 +900,12 @@ const ServicePage: FunctionComponent = () => {
   const testConnectionTab = useMemo(() => {
     return (
       <>
-        <Space className="w-full my-4 justify-between">
+        <Row className="my-4">
+          <Col span={12}>
             <AirflowMessageBanner />
-          <div>
+          </Col>
+          <Col span={12}>
+            <Space className="w-full justify-end">
               <Tooltip
                 title={
                   servicePermission.EditAll
@@ -943,8 +946,9 @@ const ServicePage: FunctionComponent = () => {
                   />
                 </Tooltip>
               )}
-          </div>
             </Space>
+          </Col>
+        </Row>
         <ServiceConnectionDetails
           connectionDetails={connectionDetails || {}}
           serviceCategory={serviceCategory}