console: introduce a generalised useNeonIntegration hook

PR-URL: https://github.com/hasura/graphql-engine-mono/pull/6042
GitOrigin-RevId: cc29154fbe7add1c9707483fd4a22b01d5e1fa13

console/src/components/Services/Data/DataSources/CreateDataSource/Neon/components/Neon/NeonBanner.stories.tsx

@@ -1,5 +1,4 @@
 import React from 'react';
-import { MdRefresh } from 'react-icons/md';
 import { ComponentMeta, Story } from '@storybook/react';
 import { within } from '@storybook/testing-library';
 import { expect } from '@storybook/jest';
@@ -13,7 +12,10 @@ export default {
 export const Base: Story = () => (
   <NeonBanner
     onClickConnect={() => window.alert('clicked connect button')}
-    status={{ status: 'default', buttonText: 'Create Neon Database for free' }}
+    status={{
+      status: 'default',
+    }}
+    buttonText="Create Neon Database for free"
   />
 );
 Base.play = async ({ canvasElement }) => {
@@ -29,8 +31,11 @@ Base.play = async ({ canvasElement }) => {
 
 export const Loading: Story = () => (
   <NeonBanner
-    onClickConnect={() => window.alert('clicked connect button')}
-    status={{ status: 'loading', buttonText: 'Authenticating' }}
+    status={{
+      status: 'loading',
+    }}
+    buttonText="Authenticating"
+    onClickConnect={() => null}
   />
 );
 Loading.play = async ({ canvasElement }) => {
@@ -45,8 +50,11 @@ Loading.play = async ({ canvasElement }) => {
 
 export const Creating: Story = () => (
   <NeonBanner
-    onClickConnect={() => window.alert('clicked connect button')}
-    status={{ status: 'loading', buttonText: 'Creating Database' }}
+    status={{
+      status: 'loading',
+    }}
+    buttonText="Creating Database"
+    onClickConnect={() => null}
   />
 );
 Creating.play = async ({ canvasElement }) => {
@@ -64,11 +72,11 @@ export const Error: Story = () => (
     onClickConnect={() => window.alert('clicked connect button')}
     status={{
       status: 'error',
-      buttonText: 'Try Again',
-      buttonIcon: <MdRefresh />,
       errorTitle: 'Error creating database',
       errorDescription: 'You have exceeded the free project limit on Neon.',
     }}
+    buttonText="Try Again"
+    icon="refresh"
   />
 );
 Error.play = async ({ canvasElement }) => {

console/src/components/Services/Data/DataSources/CreateDataSource/Neon/components/Neon/NeonBanner.tsx

@@ -1,29 +1,34 @@
-import React, { ReactElement } from 'react';
+import React from 'react';
+import { MdRefresh } from 'react-icons/md';
 import { Button } from '@/new-components/Button';
 import { IndicatorCard } from '@/new-components/IndicatorCard';
 
+const iconMap = {
+  refresh: <MdRefresh />,
+};
+
+type Status =
+  | {
+      status: 'loading';
+    }
+  | {
+      status: 'error';
+      errorTitle: string;
+      errorDescription: string;
+    }
+  | {
+      status: 'default';
+    };
+
 export type Props = {
+  status: Status;
   onClickConnect: VoidFunction;
-  status:
-    | {
-        status: 'loading';
-        buttonText: string;
-      }
-    | {
-        status: 'error';
-        buttonText: string;
-        buttonIcon: ReactElement;
-        errorTitle: string;
-        errorDescription: string;
-      }
-    | {
-        status: 'default';
-        buttonText: string;
-      };
+  buttonText: string;
+  icon?: keyof typeof iconMap;
 };
 
 export function NeonBanner(props: Props) {
-  const { status, onClickConnect } = props;
+  const { status, onClickConnect, buttonText, icon } = props;
   const isButtonDisabled = status.status === 'loading';
 
   return (
@@ -57,9 +62,9 @@ export function NeonBanner(props: Props) {
             data-testid="neon-connect-db-button"
             mode={status.status === 'loading' ? 'default' : 'primary'}
             isLoading={status.status === 'loading'}
-            loadingText={status.buttonText}
+            loadingText={buttonText}
             size="md"
-            icon={status.status === 'error' ? status.buttonIcon : undefined}
+            icon={icon ? iconMap[icon] : undefined}
             onClick={() => {
               if (!isButtonDisabled) {
                 onClickConnect();
@@ -67,7 +72,7 @@ export function NeonBanner(props: Props) {
             }}
             disabled={isButtonDisabled}
           >
-            {status.buttonText}
+            {props.buttonText}
           </Button>
         </div>
       </div>

console/src/components/Services/Data/DataSources/CreateDataSource/Neon/index.tsx

@@ -1,194 +1,110 @@
 import * as React from 'react';
-import { useDispatch } from 'react-redux';
-import { MdRefresh } from 'react-icons/md';
-import { useNeonOAuth } from './useNeonOAuth';
-import { useNeonDatabase } from './useNeonDatabase';
-import { useCreateHasuraDatasource } from './useCreateHasuraDatasource';
-import { setDBConnectionDetails } from '../../../DataActions';
+import { Dispatch } from '@/types';
 import {
   NeonBanner,
   Props as NeonBannerProps,
 } from './components/Neon/NeonBanner';
 import { getNeonDBName } from './utils';
+import { useNeonIntegration } from './useNeonIntegration';
+import _push from '../../../push';
 
-// This component creates Neon DB and calls the success/error callback
-export function Neon(props: {
-  oauthString?: string;
-  dbCreationCallback: (dbName: string) => void;
-  errorCallback: VoidFunction;
-  allDatabases: string[];
-}) {
-  const dispatch = useDispatch();
-  const {
-    state: neonDBCreationStatus,
-    create: createNeonDatabase,
-    reset: resetNeonDBCreationState,
-  } = useNeonDatabase();
+// This component deals with Neon DB creation on connect DB page
+export function Neon(props: { allDatabases: string[]; dispatch: Dispatch }) {
+  const { dispatch, allDatabases } = props;
 
-  const { startNeonOAuth, neonOauthStatus } = useNeonOAuth(props.oauthString);
+  const pushToDatasource = (dataSourceName: string) => {
+    dispatch(_push(`/data/${dataSourceName}`));
+  };
+  const pushToConnectDBPage = () => {
+    dispatch(_push(`/data/manage/connect`));
+  };
 
-  const { state: hasuraDatasourceCreationState, addHasuraDatasource } =
-    useCreateHasuraDatasource(
-      neonDBCreationStatus.status === 'success'
-        ? neonDBCreationStatus.payload.databaseUrl || ''
-        : '',
-      getNeonDBName(props.allDatabases)
-    );
+  const neonIntegrationStatus = useNeonIntegration(
+    getNeonDBName(allDatabases),
+    pushToDatasource,
+    pushToConnectDBPage,
+    dispatch
+  );
 
-  React.useEffect(() => {
-    // automatically login if creating database fails with 401 unauthorized
-    if (
-      neonDBCreationStatus.status === 'error' &&
-      neonDBCreationStatus.error === 'unauthorized' &&
-      neonOauthStatus.status === 'idle'
-    ) {
-      startNeonOAuth();
-      resetNeonDBCreationState();
-    }
-
-    // automatically create database after authentication completion
-    if (
-      neonOauthStatus.status === 'authenticated' &&
-      (neonDBCreationStatus.status === 'idle' ||
-        (neonDBCreationStatus.status === 'error' &&
-          neonDBCreationStatus.error === 'unauthorized'))
-    ) {
-      createNeonDatabase();
-    }
-  }, [neonDBCreationStatus, neonOauthStatus]);
-
-  React.useEffect(() => {
-    if (neonDBCreationStatus.status === 'success') {
-      switch (hasuraDatasourceCreationState.status) {
-        case 'idle':
-          addHasuraDatasource();
-          break;
-        case 'adding-env-var-failed':
-          dispatch(
-            setDBConnectionDetails({
-              dbURL: hasuraDatasourceCreationState.payload.dbUrl,
-              dbName: 'default',
-            })
-          );
-          props.errorCallback();
-          break;
-        case 'adding-data-source-failed':
-          dispatch(
-            setDBConnectionDetails({
-              envVar: hasuraDatasourceCreationState.payload.envVar,
-              dbName: 'default',
-            })
-          );
-          props.errorCallback();
-          break;
-        case 'success':
-          props.dbCreationCallback(
-            hasuraDatasourceCreationState.payload.dataSourceName
-          );
-          break;
-        default:
-          break;
-      }
-    }
-  }, [neonDBCreationStatus, hasuraDatasourceCreationState]);
-
-  // gets the NeonBanner render props associated with Neon DB creation
-  let neonDBCreationStatusOpts: NeonBannerProps;
-  switch (neonDBCreationStatus.status) {
+  let neonBannerProps: NeonBannerProps;
+  switch (neonIntegrationStatus.status) {
     case 'idle':
-      neonDBCreationStatusOpts = {
-        status: { status: 'default', buttonText: 'Connect Neon Database' },
-        onClickConnect: createNeonDatabase,
-      };
-      break;
-    case 'error': {
-      switch (neonDBCreationStatus.error) {
-        case 'unauthorized':
-          neonDBCreationStatusOpts = {
-            status: {
-              status: 'default',
-              buttonText: 'Connect Neon Database',
-            },
-            onClickConnect: startNeonOAuth,
-          };
-          break;
-        default:
-          neonDBCreationStatusOpts = {
-            status: {
-              status: 'error',
-              buttonText: 'Try again',
-              buttonIcon: <MdRefresh />,
-              errorTitle: 'Creating Neon Database failed',
-              errorDescription: `Error creating database: ${neonDBCreationStatus.error}`,
-            },
-            onClickConnect: createNeonDatabase,
-          };
-          break;
-      }
-      break;
-    }
-    case 'loading':
-      neonDBCreationStatusOpts = {
-        status: { status: 'loading', buttonText: 'Creating Database' },
-        onClickConnect: () => null,
-      };
-      break;
-    case 'success':
-      neonDBCreationStatusOpts = {
-        status: { status: 'loading', buttonText: 'Connecting to Hasura' },
-        onClickConnect: () => null,
-      };
-      break;
-    default:
-      // never happens; handling for placating TypeScript
-      neonDBCreationStatusOpts = {
+      neonBannerProps = {
         status: {
           status: 'default',
-          buttonText: 'Create Neon Database',
         },
-        onClickConnect: createNeonDatabase,
+        buttonText: 'Connect Neon Database',
+        onClickConnect: neonIntegrationStatus.action,
       };
       break;
-  }
-
-  // get the NeonBanner render props associated with Neon OAuth
-  let neonOAuthStatusOpts: NeonBannerProps;
-  switch (neonOauthStatus.status) {
-    case 'idle':
-      neonOAuthStatusOpts = { ...neonDBCreationStatusOpts };
-      break;
-    case 'authenticating':
-      neonOAuthStatusOpts = {
-        status: { status: 'loading', buttonText: 'Authenticating with Neon' },
+    case 'authentication-loading':
+      neonBannerProps = {
+        status: {
+          status: 'loading',
+        },
+        buttonText: 'Authenticating with Neon',
         onClickConnect: () => null,
       };
       break;
-    case 'authenticated':
-      neonOAuthStatusOpts = { ...neonDBCreationStatusOpts };
-      break;
-    case 'error':
-      neonOAuthStatusOpts = {
+    case 'authentication-error':
+      neonBannerProps = {
         status: {
           status: 'error',
-          buttonText: 'Try again',
-          buttonIcon: <MdRefresh />,
-          errorTitle: 'Error authenticating with Neon',
-          errorDescription: neonOauthStatus.error.message,
+          errorTitle: neonIntegrationStatus.title,
+          errorDescription: neonIntegrationStatus.description,
         },
-        onClickConnect: startNeonOAuth,
+        buttonText: 'Try again',
+        onClickConnect: neonIntegrationStatus.action,
+        icon: 'refresh',
+      };
+      break;
+    case 'authentication-success':
+    case 'neon-database-creation-loading':
+      neonBannerProps = {
+        status: {
+          status: 'loading',
+        },
+        buttonText: 'Creating Database',
+        onClickConnect: () => null,
+      };
+      break;
+    case 'neon-database-creation-error':
+      neonBannerProps = {
+        status: {
+          status: 'error',
+          errorTitle: neonIntegrationStatus.title,
+          errorDescription: neonIntegrationStatus.description,
+        },
+        buttonText: 'Try again',
+        onClickConnect: neonIntegrationStatus.action,
+        icon: 'refresh',
+      };
+      break;
+    case 'neon-database-creation-success':
+    case 'env-var-creation-loading':
+    case 'env-var-creation-success':
+    case 'env-var-creation-error':
+    case 'hasura-source-creation-loading':
+    case 'hasura-source-creation-error':
+    case 'hasura-source-creation-success':
+      neonBannerProps = {
+        status: {
+          status: 'loading',
+        },
+        buttonText: 'Connecting to Hasura',
+        onClickConnect: () => null,
       };
       break;
     default:
-      // never happens; handling for placating TypeScript
-      neonOAuthStatusOpts = {
+      neonBannerProps = {
         status: {
           status: 'default',
-          buttonText: 'Create Neon Database',
         },
-        onClickConnect: createNeonDatabase,
+        buttonText: 'Connect Neon Database',
+        onClickConnect: () => null,
       };
       break;
   }
 
-  return <NeonBanner {...neonOAuthStatusOpts} />;
+  return <NeonBanner {...neonBannerProps} />;
 }

console/src/components/Services/Data/DataSources/CreateDataSource/Neon/useCreateHasuraCloudDatasource.ts

@@ -1,6 +1,6 @@
 import { useState, useCallback } from 'react';
-import { useDispatch } from 'react-redux';
 import { tracingTools } from '@/features/TracingTools';
+import { Dispatch } from '@/types';
 import {
   setDBURLInEnvVars,
   verifyProjectHealthAndConnectDataSource,
@@ -44,12 +44,11 @@ type HasuraDatasourceStatus =
       payload: HasuraDBCreationPayload;
     };
 
-export function useCreateHasuraDatasource(
+export function useCreateHasuraCloudDatasource(
   dbUrl: string,
-  dataSourceName = 'default'
+  dataSourceName = 'default',
+  dispatch: Dispatch
 ) {
-  const dispatch = useDispatch();
-
   const [state, setState] = useState<HasuraDatasourceStatus>({
     status: 'idle',
   });

console/src/components/Services/Data/DataSources/CreateDataSource/Neon/useNeonIntegration.ts

@@ -0,0 +1,307 @@
+import { useEffect } from 'react';
+import { Dispatch } from '@/types';
+import { useNeonOAuth } from './useNeonOAuth';
+import { useNeonDatabase } from './useNeonDatabase';
+import { useCreateHasuraCloudDatasource } from './useCreateHasuraCloudDatasource';
+import { setDBConnectionDetails } from '../../../DataActions';
+
+type EmptyPayload = Record<string, never>;
+
+type NeonDBCreationSuccessPayload = {
+  databaseUrl: string;
+  email: string;
+};
+
+type EnvVarCreationPayload = {
+  databaseUrl: string;
+  dataSourceName: string;
+};
+type DatasourceCreationPayload = {
+  envVar: string;
+  databaseUrl: string;
+  dataSourceName: string;
+};
+
+type Idle<Status, Payload> = {
+  status: Status;
+  payload?: Payload;
+  action: VoidFunction;
+};
+
+type Loading<Status, Payload> = {
+  status: Status;
+  payload: Payload;
+};
+
+type Error<Status, Payload> = {
+  status: Status;
+  payload: Payload;
+  title: string;
+  description: string;
+  action: VoidFunction;
+};
+
+type Success<Status, Payload> = {
+  status: Status;
+  payload: Payload;
+};
+
+type NeonIntegrationStatus =
+  | Idle<'idle', EmptyPayload>
+  | Loading<'authentication-loading', EmptyPayload>
+  | Success<'authentication-success', EmptyPayload>
+  | Error<'authentication-error', EmptyPayload>
+  | Loading<'neon-database-creation-loading', EmptyPayload>
+  | Success<'neon-database-creation-success', NeonDBCreationSuccessPayload>
+  | Error<'neon-database-creation-error', EmptyPayload>
+  | Loading<'env-var-creation-loading', EnvVarCreationPayload>
+  | Success<'env-var-creation-success', DatasourceCreationPayload>
+  | Error<'env-var-creation-error', EnvVarCreationPayload>
+  | Loading<'hasura-source-creation-loading', DatasourceCreationPayload>
+  | Success<'hasura-source-creation-success', DatasourceCreationPayload>
+  | Error<'hasura-source-creation-error', DatasourceCreationPayload>;
+
+export function useNeonIntegration(
+  dataSourceName: string,
+  dbCreationCallback: (dataSourceName: string) => void, // TODO use NeonIntegrationStatus as a parameter
+  failureCallback: VoidFunction, // TODO use NeonIntegrationStatus as a parameter
+  dispatch: Dispatch
+): NeonIntegrationStatus {
+  const { startNeonOAuth, neonOauthStatus } = useNeonOAuth();
+
+  const {
+    create: createNeonDatabase,
+    state: neonDBCreationStatus,
+    reset: resetNeonDBCreationState,
+  } = useNeonDatabase();
+
+  const { state: hasuraCloudDataSourceConnectionStatus, addHasuraDatasource } =
+    useCreateHasuraCloudDatasource(
+      neonDBCreationStatus.status === 'success'
+        ? neonDBCreationStatus.payload.databaseUrl || ''
+        : '',
+      dataSourceName,
+      dispatch
+    );
+
+  useEffect(() => {
+    // automatically login if creating database fails with 401 unauthorized
+    if (
+      neonDBCreationStatus.status === 'error' &&
+      neonDBCreationStatus.error === 'unauthorized' &&
+      neonOauthStatus.status === 'idle'
+    ) {
+      startNeonOAuth();
+      resetNeonDBCreationState();
+    }
+
+    // automatically create database after authentication completion
+    if (
+      neonOauthStatus.status === 'authenticated' &&
+      (neonDBCreationStatus.status === 'idle' ||
+        (neonDBCreationStatus.status === 'error' &&
+          neonDBCreationStatus.error === 'unauthorized'))
+    ) {
+      createNeonDatabase();
+    }
+  }, [neonDBCreationStatus, neonOauthStatus]);
+
+  useEffect(() => {
+    if (neonDBCreationStatus.status === 'success') {
+      switch (hasuraCloudDataSourceConnectionStatus.status) {
+        case 'idle':
+          addHasuraDatasource();
+          break;
+        case 'adding-env-var-failed':
+          dispatch(
+            setDBConnectionDetails({
+              dbURL: hasuraCloudDataSourceConnectionStatus.payload.dbUrl,
+              dbName: 'default',
+            })
+          );
+          failureCallback();
+          break;
+        case 'adding-data-source-failed':
+          dispatch(
+            setDBConnectionDetails({
+              envVar: hasuraCloudDataSourceConnectionStatus.payload.envVar,
+              dbName: 'default',
+            })
+          );
+          failureCallback();
+          break;
+        case 'success':
+          dbCreationCallback(
+            hasuraCloudDataSourceConnectionStatus.payload.dataSourceName
+          );
+          break;
+        default:
+          break;
+      }
+    }
+  }, [neonDBCreationStatus, hasuraCloudDataSourceConnectionStatus]);
+
+  const getNeonDBCreationStatus = (): NeonIntegrationStatus => {
+    switch (neonDBCreationStatus.status) {
+      case 'idle':
+        return {
+          status: 'idle',
+          action: createNeonDatabase,
+        };
+      case 'error': {
+        switch (neonOauthStatus.status) {
+          case 'idle':
+            if (neonDBCreationStatus.error === 'unauthorized') {
+              return {
+                status: 'idle',
+                action: startNeonOAuth,
+              };
+            }
+            return {
+              status: 'neon-database-creation-error',
+              action: createNeonDatabase,
+              payload: {},
+              title: 'Error creating Neon database',
+              description: neonDBCreationStatus.error,
+            };
+
+          case 'error':
+            return {
+              status: 'authentication-error',
+              payload: {},
+              action: startNeonOAuth,
+              title: 'Error authenticating with Neon',
+              description: neonOauthStatus.error.message,
+            };
+          case 'authenticating':
+            return {
+              status: 'authentication-loading',
+              payload: {},
+            };
+          case 'authenticated':
+            return {
+              status: 'neon-database-creation-error',
+              action: createNeonDatabase,
+              payload: {},
+              title: 'Error creating Neon database',
+              description: neonDBCreationStatus.error,
+            };
+          default:
+            return {
+              status: 'idle',
+              action: startNeonOAuth,
+            };
+        }
+        break;
+      }
+      case 'loading':
+        return {
+          status: 'neon-database-creation-loading',
+          payload: {},
+        };
+      case 'success':
+        {
+          const { databaseUrl: dbUrl } = neonDBCreationStatus.payload;
+          switch (hasuraCloudDataSourceConnectionStatus.status) {
+            case 'idle':
+              return {
+                status: 'neon-database-creation-success',
+                payload: {
+                  databaseUrl: neonDBCreationStatus.payload.databaseUrl || '',
+                  email: neonDBCreationStatus.payload.email || '',
+                },
+              };
+            case 'adding-env-var':
+              return {
+                status: 'env-var-creation-loading',
+                payload: {
+                  databaseUrl: neonDBCreationStatus.payload.databaseUrl || '',
+                  dataSourceName,
+                },
+              };
+            case 'adding-env-var-failed':
+              return {
+                status: 'env-var-creation-error',
+                payload: {
+                  databaseUrl: dbUrl || '',
+                  dataSourceName,
+                },
+                action: () => null,
+                title: 'Error creating env var',
+                description:
+                  'Unexpected error adding env vars to the Hasura Cloud project',
+              };
+            case 'adding-data-source':
+              return {
+                status: 'hasura-source-creation-loading',
+                payload: {
+                  dataSourceName,
+                  envVar: hasuraCloudDataSourceConnectionStatus.payload.envVar,
+                  databaseUrl: dbUrl || '',
+                },
+              };
+            case 'success':
+              return {
+                status: 'hasura-source-creation-success',
+                payload: {
+                  databaseUrl: dbUrl || '',
+                  dataSourceName,
+                  envVar: hasuraCloudDataSourceConnectionStatus.payload.envVar,
+                },
+              };
+            case 'adding-data-source-failed':
+            default:
+              return {
+                status: 'hasura-source-creation-error',
+                payload: {
+                  databaseUrl: dbUrl || '',
+                  dataSourceName,
+                  envVar: hasuraCloudDataSourceConnectionStatus.payload.envVar,
+                },
+                action: createNeonDatabase,
+                title: 'Error creating Hasura datasource',
+                description: 'Unexpected error creating Hasura datasource',
+              };
+          }
+        }
+        break;
+      default: {
+        return {
+          status: 'idle',
+          payload: {},
+          action: createNeonDatabase,
+        };
+      }
+    }
+  };
+
+  const getNeonIntegrationStatus = (): NeonIntegrationStatus => {
+    switch (neonOauthStatus.status) {
+      case 'idle':
+        return getNeonDBCreationStatus();
+      case 'authenticating':
+        return {
+          status: 'authentication-loading',
+          payload: {},
+        };
+      case 'error':
+        return {
+          status: 'authentication-error',
+          payload: {},
+          title: 'Error authenticating with Neon',
+          description: neonOauthStatus.error.message,
+          action: startNeonOAuth,
+        };
+      case 'authenticated':
+        return getNeonDBCreationStatus();
+      default:
+        return {
+          status: 'idle',
+          payload: {},
+          action: startNeonOAuth,
+        };
+    }
+  };
+
+  return getNeonIntegrationStatus();
+}

console/src/components/Services/Data/DataSources/CreateDataSource/index.tsx

@@ -12,7 +12,6 @@ import { NotFoundError } from '../../../../Error/PageNotFound';
 import { getDataSources } from '../../../../../metadata/selector';
 import { HerokuBanner } from './Neon/components/HerokuBanner/Banner';
 import { Neon } from './Neon';
-import _push from '../../push';
 
 interface Props extends InjectedProps {}
 
@@ -38,13 +37,8 @@ const CreateDataSource: React.FC<Props> = ({
           <div className={`${styles.container} mb-md`}>
             <div className="w-full mb-md">
               <Neon
-                dbCreationCallback={dataSourceName => {
-                  dispatch(_push(`/data/${dataSourceName}`));
-                }}
-                errorCallback={() => {
-                  dispatch(_push('/data/manage/connect'));
-                }}
                 allDatabases={allDataSources.map(d => d.name)}
+                dispatch={dispatch}
               />
             </div>
             <HerokuBanner />

console/src/features/OnboardingWizard/components/NeonConnectBanner/NeonBanner.stories.tsx

@@ -1,5 +1,4 @@
 import React from 'react';
-import { MdRefresh } from 'react-icons/md';
 import { ComponentMeta, Story } from '@storybook/react';
 import { within } from '@storybook/testing-library';
 import { expect } from '@storybook/jest';
@@ -12,8 +11,9 @@ export default {
 
 export const Base: Story = () => (
   <NeonBanner
-    onButtonClick={() => window.alert('clicked connect button')}
-    status={{ status: 'default', buttonText: 'Create a Neon Database' }}
+    onClickConnect={() => window.alert('clicked connect button')}
+    status={{ status: 'default' }}
+    buttonText="Create a Neon Database"
   />
 );
 Base.play = async ({ canvasElement }) => {
@@ -31,8 +31,9 @@ Base.play = async ({ canvasElement }) => {
 
 export const Creating: Story = () => (
   <NeonBanner
-    onButtonClick={() => window.alert('clicked connect button')}
-    status={{ status: 'loading', buttonText: 'Creating Neon Database' }}
+    onClickConnect={() => window.alert('clicked connect button')}
+    status={{ status: 'loading' }}
+    buttonText="Creating Neon Database"
   />
 );
 Creating.play = async ({ canvasElement }) => {
@@ -51,14 +52,14 @@ Creating.play = async ({ canvasElement }) => {
 
 export const Error: Story = () => (
   <NeonBanner
-    onButtonClick={() => window.alert('clicked connect button')}
+    onClickConnect={() => window.alert('clicked connect button')}
     status={{
       status: 'error',
-      buttonText: 'Try Again',
-      buttonIcon: <MdRefresh />,
       errorTitle: 'Your Neon Database connection failed',
       errorDescription: 'You have exceeded the free project limit on Neon.',
     }}
+    buttonText="Try Again"
+    icon="refresh"
   />
 );
 Error.play = async ({ canvasElement }) => {

console/src/features/OnboardingWizard/components/NeonConnectBanner/NeonBanner.tsx

@@ -1,30 +1,35 @@
-import React, { ReactElement } from 'react';
+import React from 'react';
+import { MdRefresh } from 'react-icons/md';
 import { Button } from '@/new-components/Button';
 import { IndicatorCard } from '@/new-components/IndicatorCard';
 import { NeonIcon } from './NeonIcon';
 
+const iconMap = {
+  refresh: <MdRefresh />,
+};
+
+type Status =
+  | {
+      status: 'loading';
+    }
+  | {
+      status: 'error';
+      errorTitle: string;
+      errorDescription: string;
+    }
+  | {
+      status: 'default';
+    };
+
 export type Props = {
-  onButtonClick: VoidFunction;
-  status:
-    | {
-        status: 'loading';
-        buttonText: string;
-      }
-    | {
-        status: 'error';
-        buttonText: string;
-        buttonIcon: ReactElement;
-        errorTitle: string;
-        errorDescription: string;
-      }
-    | {
-        status: 'default';
-        buttonText: string;
-      };
+  status: Status;
+  onClickConnect: VoidFunction;
+  buttonText: string;
+  icon?: keyof typeof iconMap;
 };
 
 export function NeonBanner(props: Props) {
-  const { status, onButtonClick } = props;
+  const { status, onClickConnect, buttonText, icon } = props;
   const isButtonDisabled = status.status === 'loading';
 
   return (
@@ -46,17 +51,17 @@ export function NeonBanner(props: Props) {
             data-testid="onboarding-wizard-neon-connect-db-button"
             mode={status.status === 'loading' ? 'default' : 'primary'}
             isLoading={status.status === 'loading'}
-            loadingText={status.buttonText}
+            loadingText={buttonText}
             size="md"
-            icon={status.status === 'error' ? status.buttonIcon : undefined}
+            icon={icon ? iconMap[icon] : undefined}
             onClick={() => {
               if (!isButtonDisabled) {
-                onButtonClick();
+                onClickConnect();
               }
             }}
             disabled={isButtonDisabled}
           >
-            {status.buttonText}
+            {buttonText}
           </Button>
         </div>
       </div>