From ea8deb1033a4e687cf715efd81f73986caf7b3be Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 09:32:21 +0200
Subject: [PATCH 01/10] display errors

---
 .../src/components/new-diagram-form.tsx       | 17 +++++++-
 .../src/store/analysis-process.ts             |  6 +++
 .../src/store/generate-diagram-wizard.ts      | 41 +++++++++++++++++--
 3 files changed, 58 insertions(+), 6 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.tsx
index 20f8da45810..b91bf92fe2e 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.tsx
@@ -21,6 +21,7 @@ import {
   Banner,
   Button,
   css,
+  ErrorSummary,
   FormFieldContainer,
   Modal,
   ModalBody,
@@ -29,9 +30,14 @@ import {
   Option,
   Select,
   SelectTable,
+  spacing,
   TextInput,
 } from '@mongodb-js/compass-components';
 
+const footerStyles = css({
+  gap: spacing[200],
+});
+
 const FormStepContainer: React.FunctionComponent<{
   title: string;
   description?: string;
@@ -56,7 +62,7 @@ const FormStepContainer: React.FunctionComponent<{
     <>
       <ModalHeader title={title} subtitle={description}></ModalHeader>
       <ModalBody>{children}</ModalBody>
-      <ModalFooter>
+      <ModalFooter className={footerStyles}>
         <Button
           onClick={onNextClick}
           disabled={isNextDisabled}
@@ -89,6 +95,7 @@ type NewDiagramFormProps = {
   selectedDatabase: string | null;
   collections: string[];
   selectedCollections: string[];
+  error: Error | null;
 
   onCancel: () => void;
   onNameChange: (name: string) => void;
@@ -104,7 +111,7 @@ type NewDiagramFormProps = {
   onCollectionsSelectionConfirm: () => void;
 };
 
-const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
+export const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
   isModalOpen,
   formStep: currentStep,
   isLoading,
@@ -114,6 +121,7 @@ const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
   selectedDatabase,
   collections,
   selectedCollections,
+  error,
   onCancel,
   onNameChange,
   onNameConfirm,
@@ -213,6 +221,7 @@ const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
           <FormFieldContainer>
             <Select
               label=""
+              aria-label="Select connection"
               value={selectedConnectionId ?? ''}
               onChange={onConnectionSelect}
               disabled={connections.length === 0}
@@ -237,6 +246,7 @@ const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
           <FormFieldContainer>
             <Select
               label=""
+              aria-label="Select database"
               value={selectedDatabase ?? ''}
               onChange={onDatabaseSelect}
             >
@@ -311,6 +321,7 @@ const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
         isLoading={isLoading}
       >
         {formContent}
+        {error && <ErrorSummary errors={[error.message]} />}
       </FormStepContainer>
     </Modal>
   );
@@ -345,6 +356,7 @@ export default connect(
       selectedDatabase,
       databaseCollections,
       selectedCollections,
+      error,
     } = state.generateDiagramWizard;
     return {
       isModalOpen: inProgress,
@@ -360,6 +372,7 @@ export default connect(
       selectedDatabase,
       collections: databaseCollections ?? [],
       selectedCollections: selectedCollections ?? [],
+      error,
     };
   },
   {
diff --git a/packages/compass-data-modeling/src/store/analysis-process.ts b/packages/compass-data-modeling/src/store/analysis-process.ts
index e8ffc9da387..faf96e90ed0 100644
--- a/packages/compass-data-modeling/src/store/analysis-process.ts
+++ b/packages/compass-data-modeling/src/store/analysis-process.ts
@@ -209,6 +209,12 @@ export function startAnalysis(
       if (cancelController.signal.aborted) {
         dispatch({ type: AnalysisProcessActionTypes.ANALYSIS_CANCELED });
       } else {
+        services.logger.log.error(
+          services.logger.mongoLogId(1_001_000_350),
+          'DataModeling',
+          'Failed to analyze schema',
+          { err }
+        );
         dispatch({
           type: AnalysisProcessActionTypes.ANALYSIS_FAILED,
           error: err as Error,
diff --git a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
index 5f6b9faff70..c73a63f7ea1 100644
--- a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
+++ b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
@@ -23,6 +23,7 @@ export type GenerateDiagramWizardState = {
   databaseCollections: string[] | null;
   selectedCollections: string[] | null;
   automaticallyInferRelations: boolean;
+  error: Error | null;
 };
 
 export enum GenerateDiagramWizardActionTypes {
@@ -97,6 +98,7 @@ export type DatabasesFetchedAction = {
 
 export type ConnectionFailedAction = {
   type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED;
+  error: Error;
 };
 
 export type SelectDatabaseAction = {
@@ -121,6 +123,7 @@ export type CollectionsFetchedAction = {
 
 export type CollectionsFetchFailedAction = {
   type: GenerateDiagramWizardActionTypes.COLLECTIONS_FETCH_FAILED;
+  error: Error;
 };
 
 export type SelectCollectionsAction = {
@@ -151,13 +154,16 @@ export type GenerateDiagramWizardActions =
   | CollectionsFetchedAction
   | SelectCollectionsAction
   | ToggleInferRelationsAction
-  | ConfirmSelectedCollectionsAction;
+  | ConfirmSelectedCollectionsAction
+  | CollectionsFetchFailedAction
+  | ConnectionFailedAction;
 
 const INITIAL_STATE: GenerateDiagramWizardState = {
   inProgress: false,
   step: 'ENTER_NAME',
   diagramName: '',
   selectedConnectionId: null,
+  error: null,
   connectionDatabases: null,
   selectedDatabase: null,
   databaseCollections: null,
@@ -186,6 +192,7 @@ export const generateDiagramWizardReducer: Reducer<
   if (isAction(action, GenerateDiagramWizardActionTypes.CANCEL_CONFIRM_NAME)) {
     return {
       ...state,
+      error: null,
       step: 'ENTER_NAME',
     };
   }
@@ -193,6 +200,7 @@ export const generateDiagramWizardReducer: Reducer<
     return {
       ...state,
       selectedConnectionId: action.id,
+      error: null,
     };
   }
   if (
@@ -287,7 +295,8 @@ export const generateDiagramWizardReducer: Reducer<
   if (isAction(action, GenerateDiagramWizardActionTypes.CONNECTION_FAILED)) {
     return {
       ...state,
-      inProgress: false,
+      error: action.error,
+      step: 'SELECT_CONNECTION',
     };
   }
   if (
@@ -295,7 +304,8 @@ export const generateDiagramWizardReducer: Reducer<
   ) {
     return {
       ...state,
-      inProgress: false,
+      error: action.error,
+      step: 'SELECT_DATABASE',
     };
   }
   return state;
@@ -342,6 +352,13 @@ export function confirmSelectConnection(): DataModelingThunkAction<
         return;
       }
       await services.connections.connect(connectionInfo);
+      // ConnectionsService.connect does not throw an error if it fails to establish a connection,
+      // so explicitly checking if error is in the connection item and throwing it.
+      const connectionError =
+        services.connections.getConnectionById(selectedConnectionId)?.error;
+      if (connectionError) {
+        throw connectionError;
+      }
       dispatch({ type: GenerateDiagramWizardActionTypes.CONNECTION_CONNECTED });
       const mongoDBInstance =
         services.instanceManager.getMongoDBInstanceForConnection(
@@ -362,7 +379,16 @@ export function confirmSelectConnection(): DataModelingThunkAction<
           }),
       });
     } catch (err) {
-      dispatch({ type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED });
+      services.logger.log.error(
+        services.logger.mongoLogId(1_001_000_348),
+        'DataModeling',
+        'Failed to select connection',
+        { err }
+      );
+      dispatch({
+        type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED,
+        error: err as Error,
+      });
     }
   };
 }
@@ -414,8 +440,15 @@ export function confirmSelectDatabase(): DataModelingThunkAction<
           }),
       });
     } catch (err) {
+      services.logger.log.error(
+        services.logger.mongoLogId(1_001_000_349),
+        'DataModeling',
+        'Failed to select database',
+        { err }
+      );
       dispatch({
         type: GenerateDiagramWizardActionTypes.COLLECTIONS_FETCH_FAILED,
+        error: err as Error,
       });
     }
   };

From 7f0747544f7a2873a50aee6dcc3f4e5611ad63da Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 09:32:28 +0200
Subject: [PATCH 02/10] add tests

---
 .../src/components/new-diagram-form.spec.tsx  | 68 +++++++++++++++++++
 1 file changed, 68 insertions(+)
 create mode 100644 packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
new file mode 100644
index 00000000000..7e259641bc2
--- /dev/null
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -0,0 +1,68 @@
+import React from 'react';
+import type { ComponentProps } from 'react';
+import { expect } from 'chai';
+import { render, screen } from '@mongodb-js/testing-library-compass';
+import { NewDiagramForm } from './new-diagram-form';
+
+function noop() {
+  // no op
+}
+function renderNewDiagramForm(
+  props: Partial<ComponentProps<typeof NewDiagramForm>>
+) {
+  return render(
+    <NewDiagramForm
+      collections={[]}
+      databases={[]}
+      diagramName=""
+      error={null}
+      formStep="enter-name"
+      isLoading={false}
+      isModalOpen={true}
+      selectedDatabase={null}
+      selectedCollections={[]}
+      selectedConnectionId={null}
+      onCancel={noop}
+      onCollectionsSelect={noop}
+      onCollectionsSelectionConfirm={noop}
+      onConnectionConfirmSelection={noop}
+      onConnectionSelect={noop}
+      onConnectionSelectCancel={noop}
+      onDatabaseConfirmSelection={noop}
+      onDatabaseSelect={noop}
+      onDatabaseSelectCancel={noop}
+      onNameChange={noop}
+      onNameConfirm={noop}
+      onNameConfirmCancel={noop}
+      {...props}
+    />
+  );
+}
+
+describe('NewDiagramForm', function () {
+  context('it renders errors at the step', function () {
+    it('renders error when selecting a connection', function () {
+      renderNewDiagramForm({
+        formStep: 'select-connection',
+        error: new Error('Can not connect.'),
+      });
+      expect(screen.getByText('Can not connect.')).to.exist;
+    });
+
+    it('renders error when selecting a database', function () {
+      renderNewDiagramForm({
+        formStep: 'select-database',
+        error: new Error('Can not fetch databases.'),
+      });
+      expect(screen.getByText('Can not fetch databases.')).to.exist;
+    });
+
+    it('renders error when selecting collections', function () {
+      renderNewDiagramForm({
+        formStep: 'select-collections',
+        error: new Error('Can not fetch collections.'),
+      });
+      expect(screen.getByText('Can not fetch collections.')).to.exist;
+    });
+  });
+});

From 0ea0635b77629cb71a368e171c694401958251f5 Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 11:00:26 +0200
Subject: [PATCH 03/10] tests store setup

---
 packages/compass-data-modeling/src/index.ts   | 15 ++------
 .../src/provider/index.tsx                    |  2 +-
 .../compass-data-modeling/src/store/index.ts  | 36 +++++++++++++++++++
 .../src/store/reducer.ts                      | 21 ++++-------
 4 files changed, 45 insertions(+), 29 deletions(-)
 create mode 100644 packages/compass-data-modeling/src/store/index.ts

diff --git a/packages/compass-data-modeling/src/index.ts b/packages/compass-data-modeling/src/index.ts
index f58cdd6da0e..a9f5f0c5fda 100644
--- a/packages/compass-data-modeling/src/index.ts
+++ b/packages/compass-data-modeling/src/index.ts
@@ -1,5 +1,3 @@
-import { applyMiddleware, createStore } from 'redux';
-import thunk from 'redux-thunk';
 import { registerHadronPlugin } from 'hadron-app-registry';
 import { preferencesLocator } from 'compass-preferences-model/provider';
 import { connectionsLocator } from '@mongodb-js/compass-connections/provider';
@@ -7,24 +5,15 @@ import { telemetryLocator } from '@mongodb-js/compass-telemetry/provider';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import type { WorkspaceComponent } from '@mongodb-js/compass-workspaces';
 import DataModelingComponent from './components/data-modeling';
-import reducer from './store/reducer';
 import { mongoDBInstancesManagerLocator } from '@mongodb-js/compass-app-stores/provider';
 import { dataModelStorageServiceLocator } from './provider';
+import { activateDataModelingStore } from './store';
 
 const DataModelingPlugin = registerHadronPlugin(
   {
     name: 'DataModeling',
     component: DataModelingComponent,
-    activate(initialProps, services, { cleanup }) {
-      const cancelControllerRef = { current: null };
-      const store = createStore(
-        reducer,
-        applyMiddleware(
-          thunk.withExtraArgument({ ...services, cancelControllerRef })
-        )
-      );
-      return { store, deactivate: cleanup };
-    },
+    activate: activateDataModelingStore,
   },
   {
     preferences: preferencesLocator,
diff --git a/packages/compass-data-modeling/src/provider/index.tsx b/packages/compass-data-modeling/src/provider/index.tsx
index a545bbc71e2..716756627a9 100644
--- a/packages/compass-data-modeling/src/provider/index.tsx
+++ b/packages/compass-data-modeling/src/provider/index.tsx
@@ -11,7 +11,7 @@ type DataModelStorageServiceState = {
   items: MongoDBDataModelDescription[];
 };
 
-const noopDataModelStorageService: DataModelStorageServiceState &
+export const noopDataModelStorageService: DataModelStorageServiceState &
   DataModelStorage = {
   status: 'INITIAL',
   error: null,
diff --git a/packages/compass-data-modeling/src/store/index.ts b/packages/compass-data-modeling/src/store/index.ts
new file mode 100644
index 00000000000..2457a9e49de
--- /dev/null
+++ b/packages/compass-data-modeling/src/store/index.ts
@@ -0,0 +1,36 @@
+import type { PreferencesAccess } from 'compass-preferences-model/provider';
+import type { ConnectionsService } from '@mongodb-js/compass-connections/provider';
+import type { TrackFunction } from '@mongodb-js/compass-telemetry/provider';
+import type { Logger } from '@mongodb-js/compass-logging/provider';
+import type { MongoDBInstancesManager } from '@mongodb-js/compass-app-stores/provider';
+import type { DataModelStorageService } from '../provider';
+import { applyMiddleware, createStore } from 'redux';
+import reducer from './reducer';
+import thunk from 'redux-thunk';
+import type { ActivateHelpers } from 'hadron-app-registry';
+
+export type DataModelingStoreOptions = Record<string, unknown>;
+
+export type DataModelingStoreServices = {
+  preferences: PreferencesAccess;
+  connections: ConnectionsService;
+  instanceManager: MongoDBInstancesManager;
+  dataModelStorage: DataModelStorageService;
+  track: TrackFunction;
+  logger: Logger;
+};
+
+export function activateDataModelingStore(
+  _: DataModelingStoreOptions,
+  services: DataModelingStoreServices,
+  { cleanup }: ActivateHelpers
+) {
+  const cancelControllerRef = { current: null };
+  const store = createStore(
+    reducer,
+    applyMiddleware(
+      thunk.withExtraArgument({ ...services, cancelControllerRef })
+    )
+  );
+  return { store, deactivate: cleanup };
+}
diff --git a/packages/compass-data-modeling/src/store/reducer.ts b/packages/compass-data-modeling/src/store/reducer.ts
index ef499c4dbbb..dddc09fac0f 100644
--- a/packages/compass-data-modeling/src/store/reducer.ts
+++ b/packages/compass-data-modeling/src/store/reducer.ts
@@ -13,13 +13,8 @@ import { analysisProcessReducer } from './analysis-process';
 import type { DiagramActions, DiagramActionTypes } from './diagram';
 import { diagramReducer } from './diagram';
 import type { ThunkAction } from 'redux-thunk';
-import type { PreferencesAccess } from 'compass-preferences-model/provider';
-import type { ConnectionsService } from '@mongodb-js/compass-connections/provider';
-import type { TrackFunction } from '@mongodb-js/compass-telemetry/provider';
-import type { Logger } from '@mongodb-js/compass-logging/provider';
-import type { MongoDBInstancesManager } from '@mongodb-js/compass-app-stores/provider';
 import { stepReducer } from './step';
-import type { DataModelStorageService } from '../provider';
+import type { DataModelingStoreServices } from '.';
 
 const reducer = combineReducers({
   step: stepReducer,
@@ -40,18 +35,14 @@ export type DataModelingActionTypes =
 
 export type DataModelingState = ReturnType<typeof reducer>;
 
+export type DataModelingExtraArgs = DataModelingStoreServices & {
+  cancelControllerRef: { current: AbortController | null };
+};
+
 export type DataModelingThunkAction<R, A extends AnyAction> = ThunkAction<
   R,
   DataModelingState,
-  {
-    preferences: PreferencesAccess;
-    connections: ConnectionsService;
-    instanceManager: MongoDBInstancesManager;
-    dataModelStorage: DataModelStorageService;
-    track: TrackFunction;
-    logger: Logger;
-    cancelControllerRef: { current: AbortController | null };
-  },
+  DataModelingExtraArgs,
   A
 >;
 

From 9a987c4616234abf2d0914ba7fd2de75b5a5688e Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 14:26:04 +0200
Subject: [PATCH 04/10] tests

---
 .../src/components/new-diagram-form.spec.tsx  | 297 ++++++++++++++----
 .../src/components/new-diagram-form.tsx       |   6 +-
 .../src/store/generate-diagram-wizard.ts      |   8 +
 .../tests/setup-store.tsx                     | 201 ++++++++++++
 4 files changed, 455 insertions(+), 57 deletions(-)
 create mode 100644 packages/compass-data-modeling/tests/setup-store.tsx

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 7e259641bc2..a773200086a 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -1,68 +1,257 @@
 import React from 'react';
-import type { ComponentProps } from 'react';
 import { expect } from 'chai';
-import { render, screen } from '@mongodb-js/testing-library-compass';
-import { NewDiagramForm } from './new-diagram-form';
-
-function noop() {
-  // no op
-}
-function renderNewDiagramForm(
-  props: Partial<ComponentProps<typeof NewDiagramForm>>
-) {
-  return render(
-    <NewDiagramForm
-      collections={[]}
-      databases={[]}
-      diagramName=""
-      error={null}
-      formStep="enter-name"
-      isLoading={false}
-      isModalOpen={true}
-      selectedDatabase={null}
-      selectedCollections={[]}
-      selectedConnectionId={null}
-      onCancel={noop}
-      onCollectionsSelect={noop}
-      onCollectionsSelectionConfirm={noop}
-      onConnectionConfirmSelection={noop}
-      onConnectionSelect={noop}
-      onConnectionSelectCancel={noop}
-      onDatabaseConfirmSelection={noop}
-      onDatabaseSelect={noop}
-      onDatabaseSelectCancel={noop}
-      onNameChange={noop}
-      onNameConfirm={noop}
-      onNameConfirmCancel={noop}
-      {...props}
-    />
-  );
-}
+import {
+  screen,
+  userEvent,
+  waitFor,
+  within,
+} from '@mongodb-js/testing-library-compass';
+import DiagramForm from './new-diagram-form';
+import { changeName, createNewDiagram } from '../store/generate-diagram-wizard';
+import { renderWithStore } from '../../tests/setup-store';
+import type { DataModelingStore } from '../../tests/setup-store';
 
 describe('NewDiagramForm', function () {
-  context('it renders errors at the step', function () {
-    it('renders error when selecting a connection', function () {
-      renderNewDiagramForm({
-        formStep: 'select-connection',
-        error: new Error('Can not connect.'),
+  context('enter-name step', function () {
+    let store: DataModelingStore;
+    let modal: HTMLElement;
+
+    beforeEach(async () => {
+      const { store: setupStore } = await renderWithStore(<DiagramForm />);
+      store = setupStore;
+      store.dispatch(createNewDiagram());
+      modal = screen.getByTestId('new-diagram-modal');
+      expect(modal).to.be.visible;
+    });
+
+    it('allows user to enter name for the model', function () {
+      userEvent.type(
+        within(modal).getByTestId('new-diagram-name-input'),
+        'diagram-1'
+      );
+      expect(store.getState().generateDiagramWizard.diagramName).to.equal(
+        'diagram-1'
+      );
+    });
+
+    it('keeps next button disabled if diagram name is empty', function () {
+      userEvent.clear(within(modal).getByTestId('new-diagram-name-input'));
+      expect(store.getState().generateDiagramWizard.diagramName).to.equal('');
+      const button = within(modal).getByRole('button', {
+        name: /next/i,
+      });
+      expect(button.getAttribute('aria-disabled')).to.equal('true');
+    });
+
+    it('cancels process when cancel is clicked', function () {
+      userEvent.click(
+        within(modal).getByRole('button', {
+          name: /cancel/i,
+        })
+      );
+      expect(store.getState().generateDiagramWizard.inProgress).to.be.false;
+    });
+  });
+
+  context('select-connection step', function () {
+    it('shows warning if there are no connections', async function () {
+      const { store } = await renderWithStore(<DiagramForm />, {
+        connections: [],
       });
-      expect(screen.getByText('Can not connect.')).to.exist;
+      store.dispatch(createNewDiagram());
+
+      store.dispatch(changeName('diagram1'));
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /next/i,
+        })
+      );
+
+      const alert = screen.getByRole('alert');
+      expect(alert.textContent).to.contain(
+        'You do not have any connections, create a new connection first'
+      );
     });
 
-    it('renders error when selecting a database', function () {
-      renderNewDiagramForm({
-        formStep: 'select-database',
-        error: new Error('Can not fetch databases.'),
+    it('shows list of connections and allows user to select one', async function () {
+      const { store } = await renderWithStore(<DiagramForm />);
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+      expect(screen.getByText('Conn1')).to.exist;
+      expect(screen.getByText('Conn2')).to.exist;
+
+      userEvent.click(screen.getByText('Conn2'));
+      expect(store.getState().generateDiagramWizard.selectedConnectionId).to.eq(
+        'two'
+      );
+
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /next/i,
+        })
+      );
+
+      expect(store.getState().generateDiagramWizard.step).to.equal(
+        'CONNECTING'
+      );
+      await waitFor(() => {
+        expect(
+          store.getState().generateDiagramWizard.connectionDatabases
+        ).to.deep.equal(['berlin', 'sample_airbnb']);
       });
-      expect(screen.getByText('Can not fetch databases.')).to.exist;
+
+      expect(store.getState().generateDiagramWizard.step).to.equal(
+        'SELECT_DATABASE'
+      );
     });
 
-    it('renders error when selecting collections', function () {
-      renderNewDiagramForm({
-        formStep: 'select-collections',
-        error: new Error('Can not fetch collections.'),
+    // TODO
+    it.skip('shows error if it fails to connect');
+  });
+
+  context('select-database step', function () {
+    it('shows list of databases and allows user to select one', async function () {
+      const { store } = await renderWithStore(<DiagramForm />);
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      {
+        // Navigate to select db
+        userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+        userEvent.click(screen.getByText('Conn2'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+        await waitFor(() => {
+          expect(store.getState().generateDiagramWizard.step).to.equal(
+            'SELECT_DATABASE'
+          );
+        });
+      }
+
+      userEvent.click(screen.getByTestId('new-diagram-database-selector'));
+      expect(screen.getByText('berlin')).to.exist;
+      expect(screen.getByText('sample_airbnb')).to.exist;
+
+      userEvent.click(screen.getByText('sample_airbnb'));
+      expect(store.getState().generateDiagramWizard.selectedDatabase).to.eq(
+        'sample_airbnb'
+      );
+
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /next/i,
+        })
+      );
+
+      expect(store.getState().generateDiagramWizard.step).to.equal(
+        'LOADING_COLLECTIONS'
+      );
+      await waitFor(() => {
+        expect(
+          store.getState().generateDiagramWizard.selectedCollections
+        ).to.deep.equal(['listings', 'listingsAndReviews', 'reviews']);
+      });
+
+      expect(store.getState().generateDiagramWizard.step).to.equal(
+        'SELECT_COLLECTIONS'
+      );
+    });
+    // TODO
+    it.skip('shows error if it fails to fetch list of databases');
+  });
+
+  context('select-collections step', function () {
+    it('shows list of collections', async function () {
+      const { store } = await renderWithStore(<DiagramForm />);
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      {
+        // Navigate to select db
+        userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+        userEvent.click(screen.getByText('Conn2'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+        await waitFor(() => {
+          expect(store.getState().generateDiagramWizard.step).to.equal(
+            'SELECT_DATABASE'
+          );
+        });
+      }
+
+      {
+        // Navigate to select colls
+        userEvent.click(screen.getByTestId('new-diagram-database-selector'));
+        userEvent.click(screen.getByText('sample_airbnb'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+        await waitFor(() => {
+          expect(store.getState().generateDiagramWizard.step).to.equal(
+            'SELECT_COLLECTIONS'
+          );
+        });
+      }
+
+      expect(screen.getByText('listings')).to.exist;
+      expect(screen.getByText('listingsAndReviews')).to.exist;
+      expect(screen.getByText('reviews')).to.exist;
+
+      expect(store.getState().generateDiagramWizard.step).to.equal(
+        'SELECT_COLLECTIONS'
+      );
+      expect(
+        store.getState().generateDiagramWizard.selectedCollections
+      ).to.deep.equal(['listings', 'listingsAndReviews', 'reviews']);
+
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /generate/i,
+        })
+      );
+
+      await waitFor(() => {
+        expect(store.getState().generateDiagramWizard.inProgress).to.be.false;
       });
-      expect(screen.getByText('Can not fetch collections.')).to.exist;
     });
+    // TODO
+    it.skip('shows error if it fails to fetch list of collections');
   });
 });
diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.tsx
index 93dd4d8e03a..bac68927711 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.tsx
@@ -112,7 +112,7 @@ type NewDiagramFormProps = {
   onCollectionsSelectionConfirm: () => void;
 };
 
-export const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
+const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
   isModalOpen,
   formStep: currentStep,
   isLoading,
@@ -182,7 +182,7 @@ export const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
           onConfirmAction: onCollectionsSelectionConfirm,
           confirmActionLabel: 'Generate',
           isConfirmDisabled:
-            !selectedCollections || selectCollections.length === 0,
+            !selectedCollections || selectedCollections.length === 0,
           onCancelAction: onDatabaseSelectCancel,
           cancelLabel: 'Back',
         };
@@ -211,7 +211,7 @@ export const NewDiagramForm: React.FunctionComponent<NewDiagramFormProps> = ({
             <TextInput
               label="New data model name"
               value={diagramName}
-              data-testId="new-diagram-name-input"
+              data-testid="new-diagram-name-input"
               onChange={(e) => {
                 onNameChange(e.currentTarget.value);
               }}
diff --git a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
index c73a63f7ea1..fac1370cc16 100644
--- a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
+++ b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
@@ -308,6 +308,14 @@ export const generateDiagramWizardReducer: Reducer<
       step: 'SELECT_DATABASE',
     };
   }
+  if (
+    isAction(action, GenerateDiagramWizardActionTypes.CONFIRM_SELECT_DATABASE)
+  ) {
+    return {
+      ...state,
+      step: 'LOADING_COLLECTIONS',
+    };
+  }
   return state;
 };
 
diff --git a/packages/compass-data-modeling/tests/setup-store.tsx b/packages/compass-data-modeling/tests/setup-store.tsx
new file mode 100644
index 00000000000..72462d5fe8c
--- /dev/null
+++ b/packages/compass-data-modeling/tests/setup-store.tsx
@@ -0,0 +1,201 @@
+import React from 'react';
+import { renderWithConnections } from '@mongodb-js/testing-library-compass';
+import { createActivateHelpers } from 'hadron-app-registry';
+import { createNoopTrack } from '@mongodb-js/compass-telemetry/provider';
+import { createNoopLogger } from '@mongodb-js/compass-logging/provider';
+import { TestMongoDBInstanceManager } from '@mongodb-js/compass-app-stores/provider';
+import type { ConnectionInfo } from '@mongodb-js/compass-connections/provider';
+import { activateDataModelingStore } from '../src/store';
+import type { DataModelingStoreServices } from '../src/store';
+import { noopDataModelStorageService } from '../src/provider';
+import { Provider } from 'react-redux';
+import { createSandboxFromDefaultPreferences } from 'compass-preferences-model';
+
+type ConnectionInfoWithMockData = ConnectionInfo & {
+  databases: Array<{
+    _id: string;
+    name: string;
+    collectionsStatus: string;
+    collectionsLength: number;
+    collections: Array<{
+      _id: string;
+      name: string;
+      type: string;
+      sourceName: string;
+    }>;
+  }>;
+};
+
+export type DataModelingStore = Awaited<ReturnType<typeof setupStore>>;
+const testConnections = [
+  {
+    id: 'one',
+    savedConnectionType: 'favorite' as const,
+    favorite: {
+      name: 'Conn1',
+    },
+    connectionOptions: {
+      connectionString: 'mongodb://localhost:27020/test',
+    },
+    databases: [
+      {
+        _id: 'db_initial',
+        name: 'db_initial',
+        collectionsStatus: 'initial',
+        collectionsLength: 0,
+        collections: [],
+      },
+      {
+        _id: 'db_ready',
+        name: 'db_ready',
+        collectionsStatus: 'ready',
+        collectionsLength: 3,
+        collections: [
+          {
+            _id: 'db_ready.meow',
+            name: 'meow',
+            type: 'collection',
+            sourceName: '',
+            pipeline: [],
+          },
+          {
+            _id: 'db_ready.woof',
+            name: 'woof',
+            type: 'timeseries',
+            sourceName: '',
+            pipeline: [],
+          },
+          {
+            _id: 'db_ready.bwok',
+            name: 'bwok',
+            type: 'view',
+            sourceName: '',
+            pipeline: [],
+          },
+        ],
+      },
+    ],
+  },
+  {
+    id: 'two',
+    savedConnectionType: 'recent' as const,
+    favorite: {
+      name: 'Conn2',
+    },
+    connectionOptions: {
+      connectionString: 'mongodb://localhost:27021/test',
+    },
+    databases: [
+      {
+        _id: 'berlin',
+        name: 'berlin',
+        collectionsStatus: 'initial',
+        collectionsLength: 1,
+        collections: [
+          {
+            _id: 'berlin.parks',
+            name: 'parks',
+            type: 'collection',
+            sourceName: '',
+            pipeline: [],
+          },
+        ],
+      },
+      {
+        _id: 'sample_airbnb',
+        name: 'sample_airbnb',
+        collectionsStatus: 'ready',
+        collectionsLength: 3,
+        collections: [
+          {
+            _id: 'sample_airbnb.listings',
+            name: 'listings',
+            type: 'collection',
+            sourceName: '',
+            pipeline: [],
+          },
+          {
+            _id: 'sample_airbnb.reviews',
+            name: 'reviews',
+            type: 'collection',
+            sourceName: '',
+            pipeline: [],
+          },
+          {
+            _id: 'sample_airbnb.listingsAndReviews',
+            name: 'listingsAndReviews',
+            type: 'view',
+            sourceName: 'listings',
+            pipeline: [],
+          },
+        ],
+      },
+    ],
+  },
+];
+
+export const setupStore = async (
+  services: Partial<DataModelingStoreServices> = {},
+  connections: ConnectionInfoWithMockData[] = testConnections
+) => {
+  const preferences = await createSandboxFromDefaultPreferences();
+  return activateDataModelingStore(
+    {},
+    {
+      logger: createNoopLogger('TEST'),
+      track: createNoopTrack(),
+      connections: {
+        connect() {
+          return Promise.resolve();
+        },
+        getConnectionById(connId: string) {
+          const conn = connections.find((x) => x.id === connId);
+          if (!conn) {
+            return undefined;
+          }
+          return {
+            info: connections.find((x) => x.id === connId),
+          };
+        },
+        getDataServiceForConnection(connId: string) {
+          const conn = connections.find((x) => x.id === connId);
+          if (!conn) {
+            throw new Error('No available connection');
+          }
+          return {
+            listDatabases: () => {
+              return Promise.resolve(conn.databases);
+            },
+            listCollections: (database: string) => {
+              return Promise.resolve(
+                conn.databases.find((x) => x._id === database)?.collections
+              );
+            },
+          };
+        },
+      } as any,
+      preferences,
+      instanceManager: new TestMongoDBInstanceManager(),
+      dataModelStorage: noopDataModelStorageService,
+      ...services,
+    },
+    createActivateHelpers()
+  ).store;
+};
+export const renderWithStore = async (
+  component: JSX.Element,
+  {
+    services = {},
+    connections = testConnections,
+  }: {
+    services?: Partial<DataModelingStoreServices>;
+    connections?: ConnectionInfoWithMockData[];
+  } = {}
+) => {
+  const store = await setupStore(services, connections);
+  const renderResult = renderWithConnections(
+    <Provider store={store}>{component}</Provider>,
+    { connections }
+  );
+  return { ...renderResult, store };
+};

From c7f4d0de5ab56b329e28bb7e66a127d185fcd1b3 Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 14:44:42 +0200
Subject: [PATCH 05/10] clean up

---
 .../src/components/new-diagram-form.spec.tsx     | 16 ++++++++--------
 .../compass-data-modeling/tests/setup-store.tsx  | 10 ++++------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index a773200086a..5fdb7f3fd8a 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -6,7 +6,7 @@ import {
   waitFor,
   within,
 } from '@mongodb-js/testing-library-compass';
-import DiagramForm from './new-diagram-form';
+import NewDiagramForm from './new-diagram-form';
 import { changeName, createNewDiagram } from '../store/generate-diagram-wizard';
 import { renderWithStore } from '../../tests/setup-store';
 import type { DataModelingStore } from '../../tests/setup-store';
@@ -16,8 +16,8 @@ describe('NewDiagramForm', function () {
     let store: DataModelingStore;
     let modal: HTMLElement;
 
-    beforeEach(async () => {
-      const { store: setupStore } = await renderWithStore(<DiagramForm />);
+    beforeEach(() => {
+      const { store: setupStore } = renderWithStore(<NewDiagramForm />);
       store = setupStore;
       store.dispatch(createNewDiagram());
       modal = screen.getByTestId('new-diagram-modal');
@@ -54,8 +54,8 @@ describe('NewDiagramForm', function () {
   });
 
   context('select-connection step', function () {
-    it('shows warning if there are no connections', async function () {
-      const { store } = await renderWithStore(<DiagramForm />, {
+    it('shows warning if there are no connections', function () {
+      const { store } = renderWithStore(<NewDiagramForm />, {
         connections: [],
       });
       store.dispatch(createNewDiagram());
@@ -74,7 +74,7 @@ describe('NewDiagramForm', function () {
     });
 
     it('shows list of connections and allows user to select one', async function () {
-      const { store } = await renderWithStore(<DiagramForm />);
+      const { store } = renderWithStore(<NewDiagramForm />);
 
       {
         // Navigate to connections step
@@ -122,7 +122,7 @@ describe('NewDiagramForm', function () {
 
   context('select-database step', function () {
     it('shows list of databases and allows user to select one', async function () {
-      const { store } = await renderWithStore(<DiagramForm />);
+      const { store } = renderWithStore(<NewDiagramForm />);
 
       {
         // Navigate to connections step
@@ -185,7 +185,7 @@ describe('NewDiagramForm', function () {
 
   context('select-collections step', function () {
     it('shows list of collections', async function () {
-      const { store } = await renderWithStore(<DiagramForm />);
+      const { store } = renderWithStore(<NewDiagramForm />);
 
       {
         // Navigate to connections step
diff --git a/packages/compass-data-modeling/tests/setup-store.tsx b/packages/compass-data-modeling/tests/setup-store.tsx
index 72462d5fe8c..270f5874233 100644
--- a/packages/compass-data-modeling/tests/setup-store.tsx
+++ b/packages/compass-data-modeling/tests/setup-store.tsx
@@ -9,7 +9,6 @@ import { activateDataModelingStore } from '../src/store';
 import type { DataModelingStoreServices } from '../src/store';
 import { noopDataModelStorageService } from '../src/provider';
 import { Provider } from 'react-redux';
-import { createSandboxFromDefaultPreferences } from 'compass-preferences-model';
 
 type ConnectionInfoWithMockData = ConnectionInfo & {
   databases: Array<{
@@ -134,11 +133,10 @@ const testConnections = [
   },
 ];
 
-export const setupStore = async (
+export const setupStore = (
   services: Partial<DataModelingStoreServices> = {},
   connections: ConnectionInfoWithMockData[] = testConnections
 ) => {
-  const preferences = await createSandboxFromDefaultPreferences();
   return activateDataModelingStore(
     {},
     {
@@ -174,7 +172,7 @@ export const setupStore = async (
           };
         },
       } as any,
-      preferences,
+      preferences: {} as any,
       instanceManager: new TestMongoDBInstanceManager(),
       dataModelStorage: noopDataModelStorageService,
       ...services,
@@ -182,7 +180,7 @@ export const setupStore = async (
     createActivateHelpers()
   ).store;
 };
-export const renderWithStore = async (
+export const renderWithStore = (
   component: JSX.Element,
   {
     services = {},
@@ -192,7 +190,7 @@ export const renderWithStore = async (
     connections?: ConnectionInfoWithMockData[];
   } = {}
 ) => {
-  const store = await setupStore(services, connections);
+  const store = setupStore(services, connections);
   const renderResult = renderWithConnections(
     <Provider store={store}>{component}</Provider>,
     { connections }

From 7adc0bb1b05a7877863589fc158ea1ed9e453a92 Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 15:41:14 +0200
Subject: [PATCH 06/10] error tests

---
 .../src/components/new-diagram-form.spec.tsx  | 190 +++++++++++++++++-
 1 file changed, 184 insertions(+), 6 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 5fdb7f3fd8a..617384b9170 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -116,8 +116,53 @@ describe('NewDiagramForm', function () {
       );
     });
 
-    // TODO
-    it.skip('shows error if it fails to connect');
+    it('shows error if it fails to connect', async function () {
+      const { store } = renderWithStore(<NewDiagramForm />, {
+        services: {
+          connections: {
+            connect() {
+              throw new Error('Can not connect');
+            },
+            getConnectionById(id) {
+              return {
+                info: { id },
+              };
+            },
+          } as any,
+        },
+      });
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+      userEvent.click(screen.getByText('Conn2'));
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /next/i,
+        })
+      );
+
+      await waitFor(() => {
+        expect(store.getState().generateDiagramWizard.step).to.equal(
+          'SELECT_CONNECTION'
+        );
+      });
+
+      expect(store.getState().generateDiagramWizard.error?.message).to.equal(
+        'Can not connect'
+      );
+      const alert = screen.getByRole('alert');
+      expect(alert.textContent).to.contain('Can not connect');
+    });
   });
 
   context('select-database step', function () {
@@ -179,8 +224,61 @@ describe('NewDiagramForm', function () {
         'SELECT_COLLECTIONS'
       );
     });
-    // TODO
-    it.skip('shows error if it fails to fetch list of databases');
+
+    it('shows error if it fails to fetch list of databases', async function () {
+      const { store } = renderWithStore(<NewDiagramForm />, {
+        services: {
+          connections: {
+            connect() {
+              return Promise.resolve();
+            },
+            getConnectionById(id) {
+              return {
+                info: { id },
+              };
+            },
+            getDataServiceForConnection() {
+              return {
+                listDatabases() {
+                  throw new Error('Can not list databases');
+                },
+              };
+            },
+          } as any,
+        },
+      });
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+      userEvent.click(screen.getByText('Conn2'));
+      userEvent.click(
+        screen.getByRole('button', {
+          name: /next/i,
+        })
+      );
+
+      await waitFor(() => {
+        expect(store.getState().generateDiagramWizard.step).to.equal(
+          'SELECT_CONNECTION'
+        );
+      });
+
+      expect(store.getState().generateDiagramWizard.error?.message).to.equal(
+        'Can not list databases'
+      );
+      const alert = screen.getByRole('alert');
+      expect(alert.textContent).to.contain('Can not list databases');
+    });
   });
 
   context('select-collections step', function () {
@@ -251,7 +349,87 @@ describe('NewDiagramForm', function () {
         expect(store.getState().generateDiagramWizard.inProgress).to.be.false;
       });
     });
-    // TODO
-    it.skip('shows error if it fails to fetch list of collections');
+
+    it('shows error if it fails to fetch list of collections', async function () {
+      const { store } = renderWithStore(<NewDiagramForm />, {
+        services: {
+          connections: {
+            connect() {
+              return Promise.resolve();
+            },
+            getConnectionById(id) {
+              return {
+                info: { id },
+              };
+            },
+            getDataServiceForConnection() {
+              return {
+                listDatabases() {
+                  return [
+                    {
+                      _id: 'sample_airbnb',
+                      name: 'sample_airbnb',
+                    },
+                  ];
+                },
+                listCollections() {
+                  throw new Error('Can not list collections');
+                },
+              };
+            },
+          } as any,
+        },
+      });
+
+      {
+        // Navigate to connections step
+        store.dispatch(createNewDiagram());
+        store.dispatch(changeName('diagram1'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+      }
+
+      {
+        // Navigate to databases list
+        userEvent.click(screen.getByTestId('new-diagram-connection-selector'));
+        userEvent.click(screen.getByText('Conn2'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+        await waitFor(() => {
+          expect(store.getState().generateDiagramWizard.step).to.equal(
+            'SELECT_DATABASE'
+          );
+        });
+      }
+
+      {
+        // Navigate to collections
+        userEvent.click(screen.getByTestId('new-diagram-database-selector'));
+        userEvent.click(screen.getByText('sample_airbnb'));
+        userEvent.click(
+          screen.getByRole('button', {
+            name: /next/i,
+          })
+        );
+        // When it fails to load collections, we are back at SELECT_DATABASE
+        await waitFor(() => {
+          expect(store.getState().generateDiagramWizard.step).to.equal(
+            'SELECT_DATABASE'
+          );
+        });
+      }
+
+      expect(store.getState().generateDiagramWizard.error?.message).to.equal(
+        'Can not list collections'
+      );
+      const alert = screen.getByRole('alert');
+      expect(alert.textContent).to.contain('Can not list collections');
+    });
   });
 });

From e8d4312727aa3d43a87377e4389cd12d331aafc0 Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Mon, 5 May 2025 15:58:38 +0200
Subject: [PATCH 07/10] check

---
 .../src/components/new-diagram-form.spec.tsx                | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 617384b9170..54796c38b40 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -123,7 +123,7 @@ describe('NewDiagramForm', function () {
             connect() {
               throw new Error('Can not connect');
             },
-            getConnectionById(id) {
+            getConnectionById(id: string) {
               return {
                 info: { id },
               };
@@ -232,7 +232,7 @@ describe('NewDiagramForm', function () {
             connect() {
               return Promise.resolve();
             },
-            getConnectionById(id) {
+            getConnectionById(id: string) {
               return {
                 info: { id },
               };
@@ -357,7 +357,7 @@ describe('NewDiagramForm', function () {
             connect() {
               return Promise.resolve();
             },
-            getConnectionById(id) {
+            getConnectionById(id: string) {
               return {
                 info: { id },
               };

From 01f3b0383648b4e0cc354fb7a2b3887ee5b1197c Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Tue, 6 May 2025 10:48:42 +0200
Subject: [PATCH 08/10] rename folder

---
 .../src/components/new-diagram-form.spec.tsx                  | 4 ++--
 .../compass-data-modeling/{tests => test}/setup-store.tsx     | 0
 2 files changed, 2 insertions(+), 2 deletions(-)
 rename packages/compass-data-modeling/{tests => test}/setup-store.tsx (100%)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 54796c38b40..180fde7ae38 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -8,8 +8,8 @@ import {
 } from '@mongodb-js/testing-library-compass';
 import NewDiagramForm from './new-diagram-form';
 import { changeName, createNewDiagram } from '../store/generate-diagram-wizard';
-import { renderWithStore } from '../../tests/setup-store';
-import type { DataModelingStore } from '../../tests/setup-store';
+import { renderWithStore } from '../../test/setup-store';
+import type { DataModelingStore } from '../../test/setup-store';
 
 describe('NewDiagramForm', function () {
   context('enter-name step', function () {
diff --git a/packages/compass-data-modeling/tests/setup-store.tsx b/packages/compass-data-modeling/test/setup-store.tsx
similarity index 100%
rename from packages/compass-data-modeling/tests/setup-store.tsx
rename to packages/compass-data-modeling/test/setup-store.tsx

From ca1b7a87b5ce06a2ef5cb015cd42038ce591e146 Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Tue, 6 May 2025 10:53:04 +0200
Subject: [PATCH 09/10] ensure that componet is rendered

---
 .../src/components/new-diagram-form.spec.tsx                   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 180fde7ae38..003dd2ee9d9 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -157,6 +157,7 @@ describe('NewDiagramForm', function () {
         );
       });
 
+      expect(screen.getByText(/select connection/i)).to.exist;
       expect(store.getState().generateDiagramWizard.error?.message).to.equal(
         'Can not connect'
       );
@@ -273,6 +274,7 @@ describe('NewDiagramForm', function () {
         );
       });
 
+      expect(screen.getByText(/select connection/i)).to.exist;
       expect(store.getState().generateDiagramWizard.error?.message).to.equal(
         'Can not list databases'
       );
@@ -425,6 +427,7 @@ describe('NewDiagramForm', function () {
         });
       }
 
+      expect(screen.getByText(/select database/i)).to.exist;
       expect(store.getState().generateDiagramWizard.error?.message).to.equal(
         'Can not list collections'
       );

From 9f5e3a1a21ae02d521f0c2583dfcf225dbc0e48e Mon Sep 17 00:00:00 2001
From: Basit Chonka <basit.chonka@mongodb.com>
Date: Tue, 6 May 2025 11:22:52 +0200
Subject: [PATCH 10/10] handle connection error by closing modal

---
 .../src/components/new-diagram-form.spec.tsx  | 17 ++-----
 .../src/store/generate-diagram-wizard.ts      | 50 +++++++++++++++----
 2 files changed, 45 insertions(+), 22 deletions(-)

diff --git a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
index 003dd2ee9d9..571c377db0c 100644
--- a/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
+++ b/packages/compass-data-modeling/src/components/new-diagram-form.spec.tsx
@@ -152,17 +152,10 @@ describe('NewDiagramForm', function () {
       );
 
       await waitFor(() => {
-        expect(store.getState().generateDiagramWizard.step).to.equal(
-          'SELECT_CONNECTION'
-        );
+        // As it fails to connect, we are closing the modal and the toast is shown
+        // to the user with an error (which is outside of this component).
+        expect(store.getState().generateDiagramWizard.inProgress).to.be.false;
       });
-
-      expect(screen.getByText(/select connection/i)).to.exist;
-      expect(store.getState().generateDiagramWizard.error?.message).to.equal(
-        'Can not connect'
-      );
-      const alert = screen.getByRole('alert');
-      expect(alert.textContent).to.contain('Can not connect');
     });
   });
 
@@ -270,11 +263,11 @@ describe('NewDiagramForm', function () {
 
       await waitFor(() => {
         expect(store.getState().generateDiagramWizard.step).to.equal(
-          'SELECT_CONNECTION'
+          'SELECT_DATABASE'
         );
       });
 
-      expect(screen.getByText(/select connection/i)).to.exist;
+      expect(screen.getByText(/select database/i)).to.exist;
       expect(store.getState().generateDiagramWizard.error?.message).to.equal(
         'Can not list databases'
       );
diff --git a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
index fac1370cc16..9e992925577 100644
--- a/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
+++ b/packages/compass-data-modeling/src/store/generate-diagram-wizard.ts
@@ -43,6 +43,7 @@ export enum GenerateDiagramWizardActionTypes {
 
   SELECT_DATABASE = 'data-modeling/generate-diagram-wizard/SELECT_DATABASE',
   CONFIRM_SELECT_DATABASE = 'data-modeling/generate-diagram-wizard/CONFIRM_SELECT_DATABASE',
+  DATABASES_FETCH_FAILED = 'data-modeling/generate-diagram-wizard/DATABASES_FETCH_FAILED',
   COLLECTIONS_FETCHED = 'data-modeling/generate-diagram-wizard/COLLECTIONS_FETCHED',
   CANCEL_SELECTED_DATABASE = 'data-modeling/generate-diagram-wizard/CANCEL_SELECTED_DATABASE',
   COLLECTIONS_FETCH_FAILED = 'data-modeling/generate-diagram-wizard/COLLECTIONS_FETCH_FAILED',
@@ -98,7 +99,6 @@ export type DatabasesFetchedAction = {
 
 export type ConnectionFailedAction = {
   type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED;
-  error: Error;
 };
 
 export type SelectDatabaseAction = {
@@ -106,6 +106,11 @@ export type SelectDatabaseAction = {
   database: string;
 };
 
+export type DatabasesFetchFailedAction = {
+  type: GenerateDiagramWizardActionTypes.DATABASES_FETCH_FAILED;
+  error: Error;
+};
+
 export type ConfirmSelectDatabaseAction = {
   type: GenerateDiagramWizardActionTypes.CONFIRM_SELECT_DATABASE;
 };
@@ -150,6 +155,7 @@ export type GenerateDiagramWizardActions =
   | DatabasesFetchedAction
   | SelectDatabaseAction
   | ConfirmSelectDatabaseAction
+  | DatabasesFetchFailedAction
   | CancelSelectedDatabaseAction
   | CollectionsFetchedAction
   | SelectCollectionsAction
@@ -228,6 +234,15 @@ export const generateDiagramWizardReducer: Reducer<
       selectedDatabase: null,
     };
   }
+  if (
+    isAction(action, GenerateDiagramWizardActionTypes.DATABASES_FETCH_FAILED)
+  ) {
+    return {
+      ...state,
+      step: 'SELECT_DATABASE',
+      error: action.error,
+    };
+  }
   if (isAction(action, GenerateDiagramWizardActionTypes.SELECT_DATABASE)) {
     return {
       ...state,
@@ -295,8 +310,7 @@ export const generateDiagramWizardReducer: Reducer<
   if (isAction(action, GenerateDiagramWizardActionTypes.CONNECTION_FAILED)) {
     return {
       ...state,
-      error: action.error,
-      step: 'SELECT_CONNECTION',
+      inProgress: false,
     };
   }
   if (
@@ -344,16 +358,18 @@ export function confirmSelectConnection(): DataModelingThunkAction<
   | ConnectionConnectedAction
   | DatabasesFetchedAction
   | ConnectionFailedAction
+  | DatabasesFetchFailedAction
 > {
   return async (dispatch, getState, services) => {
     dispatch({
       type: GenerateDiagramWizardActionTypes.CONFIRM_SELECT_CONNECTION,
     });
+
+    const { selectedConnectionId } = getState().generateDiagramWizard;
+    if (!selectedConnectionId) {
+      return;
+    }
     try {
-      const { selectedConnectionId } = getState().generateDiagramWizard;
-      if (!selectedConnectionId) {
-        return;
-      }
       const connectionInfo =
         services.connections.getConnectionById(selectedConnectionId)?.info;
       if (!connectionInfo) {
@@ -367,6 +383,20 @@ export function confirmSelectConnection(): DataModelingThunkAction<
       if (connectionError) {
         throw connectionError;
       }
+    } catch (err) {
+      services.logger.log.error(
+        services.logger.mongoLogId(1_001_000_348),
+        'DataModeling',
+        'Failed to select connection',
+        { err }
+      );
+      dispatch({
+        type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED,
+      });
+      return;
+    }
+
+    try {
       dispatch({ type: GenerateDiagramWizardActionTypes.CONNECTION_CONNECTED });
       const mongoDBInstance =
         services.instanceManager.getMongoDBInstanceForConnection(
@@ -388,13 +418,13 @@ export function confirmSelectConnection(): DataModelingThunkAction<
       });
     } catch (err) {
       services.logger.log.error(
-        services.logger.mongoLogId(1_001_000_348),
+        services.logger.mongoLogId(1_001_000_351),
         'DataModeling',
-        'Failed to select connection',
+        'Failed to list databases',
         { err }
       );
       dispatch({
-        type: GenerateDiagramWizardActionTypes.CONNECTION_FAILED,
+        type: GenerateDiagramWizardActionTypes.DATABASES_FETCH_FAILED,
         error: err as Error,
       });
     }