diff --git a/packages/compass-data-modeling/src/components/diagram-editor.tsx b/packages/compass-data-modeling/src/components/diagram-editor.tsx
index bddf93a84ea..39a77bc9654 100644
--- a/packages/compass-data-modeling/src/components/diagram-editor.tsx
+++ b/packages/compass-data-modeling/src/components/diagram-editor.tsx
@@ -1,4 +1,4 @@
-import React, { useMemo, useState } from 'react';
+import React, { useCallback, useMemo, useRef, useState } from 'react';
 import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import {
@@ -27,6 +27,7 @@ import {
   Diagram,
   type NodeProps,
   type EdgeProps,
+  useDiagram,
 } from '@mongodb-js/diagramming';
 import type { Edit, StaticModel } from '../services/data-model-storage';
 import { UUID } from 'bson';
@@ -132,6 +133,20 @@ const DiagramEditor: React.FunctionComponent<{
   onApplyClick,
 }) => {
   const isDarkMode = useDarkMode();
+  const diagramContainerRef = useRef<HTMLDivElement | null>(null);
+  const diagram = useDiagram();
+
+  const setDiagramContainerRef = useCallback(
+    (ref: HTMLDivElement | null) => {
+      if (ref) {
+        // For debugging purposes, we attach the diagram to the ref.
+        (ref as any)._diagram = diagram;
+      }
+      diagramContainerRef.current = ref;
+    },
+    [diagram]
+  );
+
   const [applyInput, setApplyInput] = useState('{}');
 
   const isEditValid = useMemo(() => {
@@ -265,6 +280,7 @@ const DiagramEditor: React.FunctionComponent<{
   if (step === 'EDITING') {
     content = (
       <div
+        ref={setDiagramContainerRef}
         className={modelPreviewContainerStyles}
         data-testid="diagram-editor-container"
       >
diff --git a/packages/compass-e2e-tests/helpers/selectors.ts b/packages/compass-e2e-tests/helpers/selectors.ts
index f129f846383..0f66ed83304 100644
--- a/packages/compass-e2e-tests/helpers/selectors.ts
+++ b/packages/compass-e2e-tests/helpers/selectors.ts
@@ -1451,5 +1451,3 @@ export const DataModelsListItem = (diagramName: string) =>
 export const DataModelsListItemActions = (diagramName: string) =>
   `${DataModelsListItem(diagramName)} [aria-label="Show actions"]`;
 export const DataModelsListItemDeleteButton = `[data-action="delete"]`;
-export const DataModelingDiagram = '.react-flow';
-export const DataModelingDiagramNode = '.react-flow__node > div';
diff --git a/packages/compass-e2e-tests/tests/data-modeling-tab.test.ts b/packages/compass-e2e-tests/tests/data-modeling-tab.test.ts
index 3ded8ff12d0..da4361584e0 100644
--- a/packages/compass-e2e-tests/tests/data-modeling-tab.test.ts
+++ b/packages/compass-e2e-tests/tests/data-modeling-tab.test.ts
@@ -14,11 +14,23 @@ import {
   createNumbersCollection,
 } from '../helpers/insert-data';
 
+type DiagramInstance = {
+  getNodes: () => Array<{
+    id: string;
+  }>;
+};
+
 async function getDiagramNodes(browser: CompassBrowser): Promise<string[]> {
-  await browser.waitForAnimations(Selectors.DataModelingDiagram);
-  return await browser
-    .$$(Selectors.DataModelingDiagramNode)
-    .map((element) => element.getAttribute('title'));
+  const nodes = await browser.execute(function (selector) {
+    const node = document.querySelector(selector);
+    if (!node) {
+      throw new Error(`Element with selector ${selector} not found`);
+    }
+    return (
+      node as Element & { _diagram: DiagramInstance }
+    )._diagram.getNodes();
+  }, Selectors.DataModelEditor);
+  return nodes.map((x) => x.id);
 }
 
 describe('Data Modeling tab', function () {
@@ -107,6 +119,7 @@ describe('Data Modeling tab', function () {
       JSON.stringify(newModel)
     );
     await browser.clickVisible(Selectors.DataModelEditorApplyButton);
+    await browser.waitForAnimations(dataModelEditor);
 
     // Verify that the model is updated
     nodes = await getDiagramNodes(browser);
@@ -114,6 +127,7 @@ describe('Data Modeling tab', function () {
 
     // Undo the change
     await browser.clickVisible(Selectors.DataModelUndoButton);
+    await browser.waitForAnimations(dataModelEditor);
     nodes = await getDiagramNodes(browser);
     expect(nodes).to.have.lengthOf(2);
     expect(nodes).to.deep.equal([
@@ -124,6 +138,7 @@ describe('Data Modeling tab', function () {
     // Redo the change
     await browser.waitForAriaDisabled(Selectors.DataModelRedoButton, false);
     await browser.clickVisible(Selectors.DataModelRedoButton);
+    await browser.waitForAnimations(dataModelEditor);
     nodes = await getDiagramNodes(browser);
     expect(nodes).to.have.lengthOf(0);