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);
-      setCurrentWorkflow(response);
+      await deleteWorkflowById(workflowId, true);
     } 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,51 +900,55 @@ const ServicePage: FunctionComponent = () => {
   const testConnectionTab = useMemo(() => {
     return (
       <>
-        <Space className="w-full my-4 justify-between">
-          <AirflowMessageBanner />
-          <div>
-            <Tooltip
-              title={
-                servicePermission.EditAll
-                  ? t('label.edit-entity', {
-                      entity: t('label.connection'),
-                    })
-                  : t('message.no-permission-for-action')
-              }>
-              <Button
-                ghost
-                data-testid="edit-connection-button"
-                disabled={!servicePermission.EditAll}
-                type="primary"
-                onClick={goToEditConnection}>
-                {t('label.edit-entity', {
-                  entity: t('label.connection'),
-                })}
-              </Button>
-            </Tooltip>
-            {allowTestConn && isAirflowAvailable && (
+        <Row className="my-4">
+          <Col span={12}>
+            <AirflowMessageBanner />
+          </Col>
+          <Col span={12}>
+            <Space className="w-full justify-end">
               <Tooltip
                 title={
                   servicePermission.EditAll
-                    ? t('label.test-entity', {
+                    ? t('label.edit-entity', {
                         entity: t('label.connection'),
                       })
                     : t('message.no-permission-for-action')
                 }>
-                <TestConnection
-                  connectionType={serviceDetails?.serviceType ?? ''}
-                  formData={connectionDetails as ConfigData}
-                  isTestingDisabled={isTestingDisabled}
-                  serviceCategory={serviceCategory as ServiceCategory}
-                  serviceName={serviceDetails?.name}
-                  // validation is not required as we have all the data available and not in edit mode
-                  shouldValidateForm={false}
-                  showDetails={false}
-                />
+                <Button
+                  ghost
+                  data-testid="edit-connection-button"
+                  disabled={!servicePermission.EditAll}
+                  type="primary"
+                  onClick={goToEditConnection}>
+                  {t('label.edit-entity', {
+                    entity: t('label.connection'),
+                  })}
+                </Button>
               </Tooltip>
-            )}
-          </div>
-        </Space>
+              {allowTestConn && isAirflowAvailable && (
+                <Tooltip
+                  title={
+                    servicePermission.EditAll
+                      ? t('label.test-entity', {
+                          entity: t('label.connection'),
+                        })
+                      : t('message.no-permission-for-action')
+                  }>
+                  <TestConnection
+                    connectionType={serviceDetails?.serviceType ?? ''}
+                    formData={connectionDetails as ConfigData}
+                    isTestingDisabled={isTestingDisabled}
+                    serviceCategory={serviceCategory as ServiceCategory}
+                    serviceName={serviceDetails?.name}
+                    // validation is not required as we have all the data available and not in edit mode
+                    shouldValidateForm={false}
+                    showDetails={false}
+                  />
+                </Tooltip>
+              )}
+            </Space>
+          </Col>
+        </Row>
         <ServiceConnectionDetails
           connectionDetails={connectionDetails || {}}
           serviceCategory={serviceCategory}