clean up tests

Author: mabaasit

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

@@ -17,7 +17,7 @@ const DataModeling: React.FunctionComponent<DataModelingPluginInitialProps> = ({
       {showList ? (
         <SavedDiagramsList></SavedDiagramsList>
       ) : (
-        <DiagramProvider>
+        <DiagramProvider fitView>
           <DiagramEditor />
         </DiagramProvider>
       )}

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

@@ -1,4 +1,4 @@
-import React, { useEffect, useMemo, useState } from 'react';
+import React, { useMemo, useState } from 'react';
 import { connect } from 'react-redux';
 import type { DataModelingState } from '../store/reducer';
 import {
@@ -134,18 +134,6 @@ 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 {

packages/compass-e2e-tests/helpers/selectors.ts

@@ -1451,3 +1451,4 @@ export const DataModelsListItem = (diagramName: string) =>
 export const DataModelsListItemActions = (diagramName: string) =>
   `${DataModelsListItem(diagramName)} [aria-label="Show actions"]`;
 export const DataModelsListItemDeleteButton = `[data-action="delete"]`;
+export const DataModelingDiagramNode = '.react-flow__node > div';

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

@@ -14,26 +14,10 @@ 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.'
-    );
-  });
+async function getDiagramNodes(browser: CompassBrowser): Promise<string[]> {
+  return await browser
+    .$$(Selectors.DataModelingDiagramNode)
+    .map((element) => element.getAttribute('title'));
 }
 
 describe('Data Modeling tab', function () {
@@ -104,7 +88,7 @@ describe('Data Modeling tab', function () {
 
     let nodes = await getDiagramNodes(browser);
     expect(nodes).to.have.lengthOf(2);
-    expect(nodes.map((x) => x.title)).to.deep.equal([
+    expect(nodes).to.deep.equal([
       'test.testCollection1',
       'test.testCollection2',
     ]);
@@ -131,7 +115,7 @@ describe('Data Modeling tab', function () {
     await browser.clickVisible(Selectors.DataModelUndoButton);
     nodes = await getDiagramNodes(browser);
     expect(nodes).to.have.lengthOf(2);
-    expect(nodes.map((x) => x.title)).to.deep.equal([
+    expect(nodes).to.deep.equal([
       'test.testCollection1',
       'test.testCollection2',
     ]);