fix tests

Author: mabaasit

packages/compass-data-modeling/src/components/diagram-editor.tsx

@@ -1,4 +1,4 @@
-import React, { useMemo, useState } from 'react';
+import React, { useEffect, useMemo, 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';
@@ -133,6 +134,19 @@ const DiagramEditor: React.FunctionComponent<{
 }) => {
   const isDarkMode = useDarkMode();
   const [applyInput, setApplyInput] = useState('{}');
+  const diagram = useDiagram();
+
+  // For tests, we are setting the diagram instance on the window object
+  // so that we can access it in the tests.
+  useEffect(() => {
+    if (
+      process.env.APP_ENV === 'webdriverio' ||
+      process.env.NODE_ENV === 'development'
+    ) {
+      (window as any).diagramInstance = diagram;
+    }
+  }, [diagram]);
+
   const isEditValid = useMemo(() => {
     try {
       JSON.parse(applyInput);

packages/compass-e2e-tests/tests/data-modeling-tab.test.ts

@@ -14,6 +14,28 @@ import {
   createNumbersCollection,
 } from '../helpers/insert-data';
 
+type DiagramInstance = {
+  getNodes: () => Array<{ id: string; data: { title: string } }>;
+};
+
+async function getDiagramNodes(browser: CompassBrowser) {
+  return await browser.execute(() => {
+    // eslint-disable-next-line no-restricted-globals
+    if ('diagramInstance' in window) {
+      // eslint-disable-next-line no-restricted-globals
+      const diagramInstance = window.diagramInstance as DiagramInstance;
+      const nodes = diagramInstance.getNodes();
+      return nodes.map((node) => ({
+        id: node.id,
+        title: node.data.title,
+      }));
+    }
+    throw new Error(
+      'Diagram instance not found in the window object. Ensure the diagram is set for tests.'
+    );
+  });
+}
+
 describe('Data Modeling tab', function () {
   let compass: Compass;
   let browser: CompassBrowser;
@@ -80,13 +102,12 @@ describe('Data Modeling tab', function () {
     const dataModelEditor = browser.$(Selectors.DataModelEditor);
     await dataModelEditor.waitForDisplayed();
 
-    // Verify that the diagram is displayed and contains both collections
-    const text = await browser.getCodemirrorEditorText(
-      Selectors.DataModelPreview
-    );
-
-    expect(text).to.include('"ns": "test.testCollection1"');
-    expect(text).to.include('"ns": "test.testCollection2"');
+    let nodes = await getDiagramNodes(browser);
+    expect(nodes).to.have.lengthOf(2);
+    expect(nodes.map((x) => x.title)).to.deep.equal([
+      'test.testCollection1',
+      'test.testCollection2',
+    ]);
 
     // Apply change to the model
     const newModel = {
@@ -96,40 +117,30 @@ describe('Data Modeling tab', function () {
         relationships: [],
       },
     };
-    const newPreview = JSON.stringify(newModel.model, null, 2);
     await browser.setCodemirrorEditorValue(
       Selectors.DataModelApplyEditor,
       JSON.stringify(newModel)
     );
     await browser.clickVisible(Selectors.DataModelEditorApplyButton);
 
     // Verify that the model is updated
-    const updatedText = await browser.getCodemirrorEditorText(
-      Selectors.DataModelPreview
-    );
-    expect(updatedText).to.equal(newPreview);
+    nodes = await getDiagramNodes(browser);
+    expect(nodes).to.have.lengthOf(0);
 
     // Undo the change
     await browser.clickVisible(Selectors.DataModelUndoButton);
-    await browser.waitUntil(async () => {
-      const textAfterUndo = await browser.getCodemirrorEditorText(
-        Selectors.DataModelPreview
-      );
-      return (
-        textAfterUndo.includes('"ns": "test.testCollection1"') &&
-        textAfterUndo.includes('"ns": "test.testCollection2"')
-      );
-    });
+    nodes = await getDiagramNodes(browser);
+    expect(nodes).to.have.lengthOf(2);
+    expect(nodes.map((x) => x.title)).to.deep.equal([
+      'test.testCollection1',
+      'test.testCollection2',
+    ]);
 
     // Redo the change
     await browser.waitForAriaDisabled(Selectors.DataModelRedoButton, false);
     await browser.clickVisible(Selectors.DataModelRedoButton);
-    await browser.waitUntil(async () => {
-      const redoneText = await browser.getCodemirrorEditorText(
-        Selectors.DataModelPreview
-      );
-      return redoneText === newPreview;
-    });
+    nodes = await getDiagramNodes(browser);
+    expect(nodes).to.have.lengthOf(0);
 
     // Open a new tab
     await browser.openNewTab();
@@ -139,10 +150,8 @@ describe('Data Modeling tab', function () {
     await browser.$(Selectors.DataModelEditor).waitForDisplayed();
 
     // Verify that the diagram has the latest changes
-    const savedText = await browser.getCodemirrorEditorText(
-      Selectors.DataModelPreview
-    );
-    expect(savedText).to.equal(newPreview);
+    nodes = await getDiagramNodes(browser);
+    expect(nodes).to.have.lengthOf(0);
 
     // Open a new tab
     await browser.openNewTab();