diff --git a/package.json b/package.json
index f9c2892..7016ae1 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,9 @@
 {
     "name": "ui-tests",
     "version": "0.0.1",
+    "exports": {
+        "./e2e-utils": "./tests/e2e/utils.ts"
+    },
     "description": "",
     "homepage": "https://github.com/Pometry/ui-tests",
     "bugs": {
diff --git a/tests/e2e/graph.spec.ts b/tests/e2e/graph.spec.ts
index 653580f..940c9e6 100644
--- a/tests/e2e/graph.spec.ts
+++ b/tests/e2e/graph.spec.ts
@@ -1,4 +1,4 @@
-import { expect, Page } from '@playwright/test';
+import { expect } from '@playwright/test';
 import { test } from '../fixtures';
 import {
     changeTab,
@@ -8,6 +8,7 @@ import {
     doubleClickOnNode,
     dragSlider,
     fillInStyling,
+    fitView,
     getGraphState,
     navigateToGraphPageBySearch,
     navigateToSavedGraphBySavedGraphsTable,
@@ -16,13 +17,6 @@ import {
     waitForLayoutToFinish,
 } from './utils';
 
-async function fitView(page: Page) {
-    await page
-        .getByRole('button', { name: 'Fit all nodes within visible region' })
-        .click();
-    await waitForLayoutToFinish(page);
-}
-
 test('Close right hand side panel button and open again', async ({ page }) => {
     await page.goto('/graph/vanilla/event?initialNodes=%5B%5D');
 
diff --git a/tests/e2e/utils.ts b/tests/e2e/utils.ts
index eb4a150..698bc60 100644
--- a/tests/e2e/utils.ts
+++ b/tests/e2e/utils.ts
@@ -23,6 +23,13 @@ type BrowserWindow = Window & {
     };
 };
 
+export async function fitView(page: Page) {
+    await page
+        .getByRole('button', { name: 'Fit all nodes within visible region' })
+        .click();
+    await waitForLayoutToFinish(page);
+}
+
 async function getNodePosition(
     page: Page,
     displayName: string,
@@ -377,6 +384,7 @@ export async function fillInStyling(
 }
 
 interface GraphState {
+    highlighted: { id: string }[];
     selected: string[];
     nodes: {
         id: string;
@@ -386,19 +394,31 @@ interface GraphState {
 }
 
 export async function getGraphState(page: Page): Promise<GraphState> {
-    return page.evaluate(() => {
-        const graph = (window as BrowserWindow).__G6_GRAPH__;
-        if (!graph) throw new Error('__G6_GRAPH__ not found on window');
-        const data = graph.getData();
-        return {
-            selected: data.nodes
-                .filter((n) => n.states?.includes('selected'))
-                .map((n) => n.id),
-            nodes: data.nodes.map((n) => ({
-                id: n.id,
-                colour: n.style?.fill,
-                size: n.style?.size,
-            })),
-        };
-    });
+    const handle = await page.waitForFunction(
+        () => {
+            const graph = (window as BrowserWindow).__G6_GRAPH__;
+            if (!graph) return undefined;
+            const data = graph.getData();
+            const anyDisabled = data.nodes.some((n) =>
+                n.states?.includes('disabled'),
+            );
+            return {
+                highlighted: anyDisabled
+                    ? data.nodes
+                          .filter((n) => !n.states?.includes('disabled'))
+                          .map((n) => ({ id: n.id }))
+                    : [],
+                selected: data.nodes
+                    .filter((n) => n.states?.includes('selected'))
+                    .map((n) => n.id),
+                nodes: data.nodes.map((n) => ({
+                    id: n.id,
+                    colour: n.style?.fill,
+                    size: n.style?.size,
+                })),
+            };
+        },
+        { timeout: 10000 },
+    );
+    return handle.jsonValue() as Promise<GraphState>;
 }