refactor register connection to use hook

Author: m0nggh

packages/frontend/src/components/ChooseConnectionSubstep/index.tsx

@@ -63,16 +63,10 @@ function ChooseConnectionSubstep(
       connectionId: connection?.id,
       flowId: supportsConnectionRegistration ? step.flowId : undefined,
     },
-    // cache-first to prevent the test connection from being called multiple times
-    fetchPolicy: 'cache-first',
     skip: !connection?.id,
   })
 
   const isTestStepValid = useMemo(() => {
-    if (application.key === APP_ALLOWING_EMPTY_CONNECTION) {
-      return true
-    }
-
     if (testResultLoading || !testConnectionData?.testConnection) {
       return null
     }
@@ -83,13 +77,20 @@ function ChooseConnectionSubstep(
       return false
     }
     return true
-  }, [application.key, testConnectionData, testResultLoading])
+  }, [testConnectionData, testResultLoading])
 
   const connectionStatus: ConnectionStatus = useMemo(() => {
     if (!connection) {
-      return {
-        text: 'Not connected',
-        color: 'yellow.200',
+      if (application.key === APP_ALLOWING_EMPTY_CONNECTION) {
+        return {
+          text: 'No connection',
+          color: 'green.500',
+        }
+      } else {
+        return {
+          text: 'Not connected',
+          color: 'yellow.200',
+        }
       }
     } else if (testResultLoading) {
       return {

packages/frontend/src/components/FlowStepConfigurationModal/ChooseAndAddConnection/AddConnection.tsx

@@ -20,24 +20,29 @@ import Form from '../../Form'
 import { infoboxMdComponents } from '../../MarkdownRenderer/CustomMarkdownComponents'
 import { DEFAULT_ADD_CONNECTION_LABEL } from '../constants'
 import { FlowStepConfigurationContext } from '../FlowStepConfigurationContext'
+import { useConnectionVerification } from '../hooks/useConnectionRegistration'
 import InvalidModalScreen from '../InvalidModalScreen'
 
 import ConnectionHeader from './ConnectionHeader'
 
 type AddConnectionProps = {
-  onSubmit: (response: Record<string, unknown>) => void
-  onBack: () => void
+  handleConnectionChange: (
+    connectionId: string,
+    shouldRefetch: boolean,
+  ) => Promise<void>
+  onCreateOrUpdateStep: (connectionId: string) => Promise<void>
 }
 
 type Response = {
   [key: string]: any
 }
 
 export default function AddConnection(props: AddConnectionProps): JSX.Element {
-  const { onSubmit, onBack } = props
-  const {
-    modalState: { selectedApp },
-  } = useContext(FlowStepConfigurationContext)
+  const { handleConnectionChange, onCreateOrUpdateStep } = props
+  const { modalState, patchModalState } = useContext(
+    FlowStepConfigurationContext,
+  )
+  const { selectedApp } = modalState
   const { name, authDocUrl, key, auth } = selectedApp || {}
 
   const [error, setError] = useState<IJSONObject | null>(null)
@@ -61,6 +66,57 @@ export default function AddConnection(props: AddConnectionProps): JSX.Element {
     }
   }, [])
 
+  const supportsConnectionRegistration =
+    !!selectedApp?.auth?.connectionRegistrationType
+
+  const { onRegisterConnection, testConnection } = useConnectionVerification({
+    supportsConnectionRegistration,
+  })
+
+  /**
+   * Test connection first for apps that support connection registration
+   * If connection is not verified, redirect to choose-connection screen to connect
+   * If connection is verified, register connection and create step
+   *
+   * For apps without connection registration, create step directly
+   */
+  const handleAddConnection = useCallback(
+    async (response: Record<string, any>) => {
+      const newConnectionId = response?.createConnection?.id as
+        | string
+        | undefined
+
+      if (newConnectionId) {
+        patchModalState({ isLoading: true })
+        if (supportsConnectionRegistration) {
+          const testResult = await testConnection(newConnectionId)
+          // If connection is not verified and has an error message
+          if (!testResult?.registrationVerified && !!testResult?.message) {
+            await handleConnectionChange(newConnectionId, true)
+            patchModalState({
+              selectedConnectionId: newConnectionId,
+              currentScreen: 'choose-connection',
+              isLoading: false,
+            })
+            return
+          } else {
+            await onRegisterConnection(newConnectionId)
+          }
+        }
+        // Create or update step at the end
+        await onCreateOrUpdateStep(newConnectionId)
+      }
+    },
+    [
+      handleConnectionChange,
+      onCreateOrUpdateStep,
+      onRegisterConnection,
+      patchModalState,
+      supportsConnectionRegistration,
+      testConnection,
+    ],
+  )
+
   const submitHandler: SubmitHandler<FieldValues> = useCallback(
     async (data) => {
       if (!steps) {
@@ -95,15 +151,17 @@ export default function AddConnection(props: AddConnectionProps): JSX.Element {
         stepIndex++
 
         if (stepIndex === steps.length) {
-          onSubmit(response)
+          await handleAddConnection(response)
         }
       }
 
       setInProgress(false)
     },
-    [key, steps, onSubmit],
+    [key, steps, handleAddConnection],
   )
 
+  const onBack = () => patchModalState({ currentScreen: 'choose-connection' })
+
   if (!selectedApp) {
     return <InvalidModalScreen />
   }
@@ -129,7 +187,7 @@ export default function AddConnection(props: AddConnectionProps): JSX.Element {
           DEFAULT_ADD_CONNECTION_LABEL
         }
       />
-      <ModalCloseButton mt={6} size="xs" />
+      <ModalCloseButton mt={2} size="xs" />
 
       <ModalBody mt={2}>
         {auth?.connectionType !== 'user-added' ? (

packages/frontend/src/components/FlowStepConfigurationModal/ChooseAndAddConnection/ChooseConnection.tsx

@@ -1,14 +1,9 @@
-import type { ITestConnectionOutput } from '@plumber/types'
-
-import { useCallback, useContext } from 'react'
+import { useContext } from 'react'
 import { BiChevronLeft } from 'react-icons/bi'
-import { useMutation, useQuery } from '@apollo/client'
 import { Flex, ModalBody, ModalHeader } from '@chakra-ui/react'
 import { Button, ModalCloseButton } from '@opengovsg/design-system-react'
 
 import { EditorContext } from '@/contexts/Editor'
-import { REGISTER_CONNECTION } from '@/graphql/mutations/register-connection'
-import { TEST_CONNECTION } from '@/graphql/queries/test-connection'
 
 import { DEFAULT_CHOOSE_CONNECTION_LABEL } from '../constants'
 import { FlowStepConfigurationContext } from '../FlowStepConfigurationContext'
@@ -22,8 +17,11 @@ import { ConnectionDropdownOption } from '.'
 interface ChooseConnectionProps {
   appConnectionsLoading: boolean
   connectionOptions: ConnectionDropdownOption[]
-  handleConnectionChange: (value: string, shouldRefetch: boolean) => void
-  onCreateOrUpdateStep: (connectionId: string) => void
+  handleConnectionChange: (
+    value: string,
+    shouldRefetch: boolean,
+  ) => Promise<void>
+  onCreateOrUpdateStep: (connectionId: string) => Promise<void>
 }
 
 export default function ChooseConnection(
@@ -35,55 +33,14 @@ export default function ChooseConnection(
     handleConnectionChange,
     onCreateOrUpdateStep,
   } = props
-  const { readOnly, flowId } = useContext(EditorContext)
+  const { readOnly } = useContext(EditorContext)
 
   const { modalState, patchModalState, step } = useContext(
     FlowStepConfigurationContext,
   )
   const { selectedApp, selectedConnectionId } = modalState
-
-  const supportsConnectionRegistration =
-    !!selectedApp?.auth?.connectionRegistrationType
-
   const connectionModalLabel = selectedApp?.auth?.connectionModalLabel
 
-  const {
-    loading: testResultLoading,
-    refetch: retestConnection,
-    data: testConnectionData,
-  } = useQuery<{
-    testConnection: ITestConnectionOutput
-  }>(TEST_CONNECTION, {
-    variables: {
-      connectionId: selectedConnectionId,
-      flowId: supportsConnectionRegistration ? flowId : undefined,
-    },
-    skip: !selectedConnectionId,
-  })
-
-  const [registerConnection, { loading: registerConnectionLoading }] =
-    useMutation(REGISTER_CONNECTION)
-
-  const onRegisterConnection = useCallback(async () => {
-    if (selectedConnectionId && supportsConnectionRegistration) {
-      await registerConnection({
-        variables: {
-          input: {
-            connectionId: selectedConnectionId,
-            flowId,
-          },
-        },
-      })
-      await retestConnection()
-    }
-  }, [
-    selectedConnectionId,
-    supportsConnectionRegistration,
-    registerConnection,
-    flowId,
-    retestConnection,
-  ])
-
   const onBack = () => {
     if (!selectedApp) {
       return
@@ -133,7 +90,7 @@ export default function ChooseConnection(
           }
         />
       </ModalHeader>
-      <ModalCloseButton mt={6} size="xs" />
+      <ModalCloseButton mt={2} size="xs" />
 
       <ModalBody>
         <Flex flexDir="column" alignItems="center" gap={6}>
@@ -149,13 +106,10 @@ export default function ChooseConnection(
               }
             />
             <SetConnectionButton
-              onNextStep={() => onCreateOrUpdateStep(selectedConnectionId)}
-              onRegisterConnection={onRegisterConnection}
+              onNextStep={async () =>
+                await onCreateOrUpdateStep(selectedConnectionId)
+              }
               readOnly={readOnly}
-              supportsConnectionRegistration={supportsConnectionRegistration}
-              testResult={testConnectionData?.testConnection}
-              testResultLoading={testResultLoading}
-              registerConnectionLoading={registerConnectionLoading}
               isNewStep={!step?.appKey || !step?.key}
             />
           </Flex>