Simplify auth

Author: TheNando

cmd/ui/tests/README.md

@@ -24,7 +24,7 @@ A Playwright **setup project** that runs once per Playwright invocation before a
 2. Delegates to `bh-playwright-testing` helpers to manage per-theme auth sessions
 
 > About auth session management:\
-> `loginAndSnapshotThemes` logs in via `/ui/login`, snapshots `storageState` to `playwright/.auth/user-light.json`, toggles dark mode (waiting for the throttled `persistedState` write to land), then snapshots `playwright/.auth/user-dark.json`. Capturing both theme snapshots from a single session avoids the parallel-login race where two setups as the same user would invalidate each other's session. Downstream projects load the snapshot that matches their theme via `authStorageStateFor(theme)`.
+> `loginAndSnapshotThemes` logs in via `/ui/login`, then writes the canonical light/dark `persistedState` (carrying over the session token from login) directly into localStorage to snapshot `storageState` to `playwright/.auth/user-light.json` and `playwright/.auth/user-dark.json`. Capturing both theme snapshots from a single session avoids the parallel-login race where two setups as the same user would invalidate each other's session. Downstream projects load the snapshot that matches their theme via `authStorageStateFor(theme)`.
 
 ## Conventions
 

cmd/ui/tests/a11y/download-collectors.a11y.spec.ts

@@ -23,8 +23,8 @@ test.describe('Download Collectors page accessibility', () => {
     }, testInfo) => {
         await page.goto('/ui/download-collectors');
 
-        // Wait for page label to load
-        await page.getByRole('heading', { name: 'Download Collectors' }).waitFor({ state: 'visible' });
+        // Wait for collector download button to appear
+        await page.getByRole('button', { name: 'Download SharpHound' }).waitFor({ state: 'visible' });
 
         const results = await makeAxeBuilder().include('#content-wrapper').analyze();
         await expectNoAccessibilityViolations(testInfo, results, { page });

cmd/ui/tests/global.setup.ts

@@ -14,7 +14,7 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
-import { expect, test as setup } from 'bh-playwright-testing';
+import { test as setup } from 'bh-playwright-testing';
 import { loginAndSnapshotThemes } from 'bh-playwright-testing/auth';
 import { installGraphHasDataStub } from 'bh-playwright-testing/stubs/graph-has-data';
 import { authStorageStateFor, type Theme } from 'bh-playwright-testing/themes';
@@ -41,15 +41,5 @@ setup('Generate and cache auth state for light and dark theme', async ({ page })
         username,
         password,
         storageStatePathFor: (theme: Theme) => path.resolve(__dirname, '..', authStorageStateFor(theme)),
-        // React Query still goes through one undefined render before the stub lands, and MUI's
-        // Dialog backdrop persists during its close transition. If the dialog raced ahead of us,
-        // dismiss it so the dark-mode click is not intercepted.
-        dismissPostLogin: async (page) => {
-            const cancel = page.getByTestId('confirmation-dialog_button-no');
-            await cancel.click({ timeout: 2_000 }).catch(() => {
-                /* dialog never opened — nothing to dismiss */
-            });
-            await expect(cancel).toBeHidden();
-        },
     });
 });

packages/javascript/bh-playwright-testing/src/auth.ts

@@ -18,27 +18,44 @@ import type { Page } from '@playwright/test';
 import { expect } from '@playwright/test';
 import type { Theme } from './themes';
 
+const persistedState = {
+    auth: {
+        sessionToken: '',
+    },
+    global: {
+        view: {
+            darkMode: false,
+            autoRunQueries: true,
+            notifications: [],
+            exploreLayout: 'sequential',
+            isExploreTableSelected: false,
+            isExploreLayoutSelected: false,
+            selectedExploreTableColumns: { kind: true, label: true, objectId: true, isTierZero: true },
+            pinnedExploreTableColumns: ['action-menu', 'kind', 'label'],
+            timeoutSetting: false,
+            isExploreGraphHighlight: true,
+            isExploreGraphLabelClip: true,
+        },
+    },
+};
+
 export type LoginAndSnapshotThemesOptions = {
     page: Page;
     username: string;
     password: string;
     // Resolves a per-theme `storageState` file path. Callers typically wrap
     // `authStorageStateFor(theme)` from `./themes` with `path.resolve(...)`.
     storageStatePathFor: (theme: Theme) => string;
-    // Optional hook for dismissing any post-login overlay (e.g. an upload dialog) that
-    // could intercept the dark-mode toggle click. Invoked after the login form clears.
-    dismissPostLogin?: (page: Page) => Promise<void>;
 };
 
 // Login once and snapshot `storageState` for both light and dark themes.
 //
 // Capturing both snapshots from a single session avoids the parallel-login race where two
-// setups as the same user would invalidate each other's session. This relies on the BH
-// shared UI shell (login form labels "Email Address" / "Password", the LOGIN submit button,
-// the `global_nav-dark-mode` toggle, and the `persistedState` localStorage key written by
-// the global store's throttled subscriber) being identical across consumers.
+// setups as the same user would invalidate each other's session. Rather than toggling the
+// theme through the UI, we write the canonical light/dark `persistedState` directly into
+// localStorage (carrying over the real session token from login) before each snapshot.
 export async function loginAndSnapshotThemes(opts: LoginAndSnapshotThemesOptions): Promise<void> {
-    const { page, username, password, storageStatePathFor, dismissPostLogin } = opts;
+    const { page, username, password, storageStatePathFor } = opts;
 
     try {
         await page.goto('/ui/login');
@@ -48,33 +65,27 @@ export async function loginAndSnapshotThemes(opts: LoginAndSnapshotThemesOptions
         // Use `exact` as some environments also have "Login Via SSO".
         await page.getByRole('button', { name: 'LOGIN', exact: true }).click();
 
-        // Race the success path (navigation off /ui/login) against the most common failure
-        // path (rejected creds leave you on /ui/login with an inline toast). This surfaces
-        // auth failures in ~15s with a clear message.
+        // Rejected creds leave you on /ui/login, so waiting to navigate away surfaces auth
+        // failures in ~15s with a clear message.
         await expect(page).not.toHaveURL(/\/ui\/login(\?|$)/, { timeout: 15_000 });
-        await page.getByTestId('global_nav-dark-mode').waitFor({ state: 'visible' });
-
-        if (dismissPostLogin) {
-            await dismissPostLogin(page);
-        }
 
-        // Snapshot light first, while the UI is still in its default theme.
-        await page.context().storageState({ path: storageStatePathFor('light') });
+        // Carry over the real session token written by the login flow so the snapshots stay
+        // authenticated after we overwrite `persistedState`.
+        const sessionToken = await page.evaluate(() => {
+            const raw = localStorage.getItem('persistedState');
+            return raw ? JSON.parse(raw)?.auth?.sessionToken ?? '' : '';
+        });
 
-        // Click the nav dark-mode toggle. The inner Switch is `inert`, so the click handler
-        // lives on the parent item targeted by this test id.
-        await page.getByTestId('global_nav-dark-mode').click();
-
-        // The store persists via a throttled (1s) subscriber. Poll until the new value lands
-        // in localStorage before snapshotting storageState, so the next test boots in dark mode.
-        await expect
-            .poll(async () => {
-                const raw = await page.evaluate(() => localStorage.getItem('persistedState'));
-                return raw ? JSON.parse(raw)?.global?.view?.darkMode : null;
-            })
-            .toBe(true);
-
-        await page.context().storageState({ path: storageStatePathFor('dark') });
+        for (const [theme, darkMode] of [
+            ['light', false],
+            ['dark', true],
+        ] as const) {
+            const state = structuredClone(persistedState);
+            state.global.view.darkMode = darkMode;
+            state.auth.sessionToken = sessionToken;
+            await page.evaluate((next) => localStorage.setItem('persistedState', JSON.stringify(next)), state);
+            await page.context().storageState({ path: storageStatePathFor(theme) });
+        }
     } catch (error) {
         throw new Error(`Auth setup failed at ${page.url()}: ${(error as Error).message}`, { cause: error });
     }