Stabilize company selection bootstrap tests

ui/src/context/CompanyContext.test.tsx

@@ -1,114 +1,43 @@
-// @vitest-environment jsdom
+import { describe, expect, it } from "vitest";
+import { resolveBootstrapCompanySelection } from "./CompanyContext";
 
-import { act, useEffect } from "react";
-import { createRoot, type Root } from "react-dom/client";
-import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { CompanyProvider, useCompany } from "./CompanyContext";
+const activeCompany = { id: "company-1" };
+const archivedCompany = { id: "archived-company" };
 
-const mockCompaniesApi = vi.hoisted(() => ({
-  list: vi.fn(),
-  create: vi.fn(),
-}));
-
-vi.mock("../api/companies", () => ({
-  companiesApi: mockCompaniesApi,
-}));
-
-function Probe({ onSelectedCompanyId }: { onSelectedCompanyId: (companyId: string | null) => void }) {
-  const { selectedCompanyId } = useCompany();
-  useEffect(() => {
-    onSelectedCompanyId(selectedCompanyId);
-  }, [onSelectedCompanyId, selectedCompanyId]);
-  return <div data-selected-company-id={selectedCompanyId ?? ""} />;
-}
-
-describe("CompanyProvider", () => {
-  let container: HTMLDivElement;
-  let root: Root;
-  let queryClient: QueryClient;
-
-  beforeEach(() => {
-    (globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }).IS_REACT_ACT_ENVIRONMENT = true;
-    localStorage.clear();
-    container = document.createElement("div");
-    document.body.appendChild(container);
-    root = createRoot(container);
-    queryClient = new QueryClient({
-      defaultOptions: {
-        queries: { retry: false },
-      },
-    });
+describe("resolveBootstrapCompanySelection", () => {
+  it("does not expose a stale stored company id before companies load", () => {
+    expect(resolveBootstrapCompanySelection({
+      companies: [],
+      sidebarCompanies: [],
+      selectedCompanyId: null,
+      storedCompanyId: "stale-company",
+    })).toBeNull();
   });
 
-  afterEach(() => {
-    act(() => {
-      root.unmount();
-    });
-    queryClient.clear();
-    container.remove();
-    vi.clearAllMocks();
+  it("replaces a stale stored company id with the first loaded company", () => {
+    expect(resolveBootstrapCompanySelection({
+      companies: [activeCompany],
+      sidebarCompanies: [activeCompany],
+      selectedCompanyId: null,
+      storedCompanyId: "stale-company",
+    })).toBe("company-1");
   });
 
-  it("does not expose a stale stored company id before companies load", async () => {
-    localStorage.setItem("paperclip.selectedCompanyId", "stale-company");
-    mockCompaniesApi.list.mockImplementation(() => new Promise(() => {}));
-    const seen: Array<string | null> = [];
-
-    await act(async () => {
-      root.render(
-        <QueryClientProvider client={queryClient}>
-          <CompanyProvider>
-            <Probe onSelectedCompanyId={(companyId) => seen.push(companyId)} />
-          </CompanyProvider>
-        </QueryClientProvider>,
-      );
-    });
-
-    expect(seen).toEqual([null]);
+  it("keeps a valid selected company ahead of stored bootstrap state", () => {
+    expect(resolveBootstrapCompanySelection({
+      companies: [activeCompany],
+      sidebarCompanies: [activeCompany],
+      selectedCompanyId: "company-1",
+      storedCompanyId: "stale-company",
+    })).toBe("company-1");
   });
 
-  it("replaces a stale stored company id with the first loaded company", async () => {
-    localStorage.setItem("paperclip.selectedCompanyId", "stale-company");
-    mockCompaniesApi.list.mockResolvedValue([
-      {
-        id: "company-1",
-        name: "Paperclip",
-        description: null,
-        status: "active",
-        issuePrefix: "PAP",
-        issueCounter: 1,
-        budgetMonthlyCents: 0,
-        spentMonthlyCents: 0,
-        requireBoardApprovalForNewAgents: false,
-        feedbackDataSharingEnabled: false,
-        feedbackDataSharingConsentAt: null,
-        feedbackDataSharingConsentByUserId: null,
-        feedbackDataSharingTermsVersion: null,
-        brandColor: null,
-        logoAssetId: null,
-        logoUrl: null,
-        createdAt: new Date(),
-        updatedAt: new Date(),
-      },
-    ]);
-    const seen: Array<string | null> = [];
-
-    await act(async () => {
-      root.render(
-        <QueryClientProvider client={queryClient}>
-          <CompanyProvider>
-            <Probe onSelectedCompanyId={(companyId) => seen.push(companyId)} />
-          </CompanyProvider>
-        </QueryClientProvider>,
-      );
-    });
-
-    await act(async () => {
-      await vi.waitFor(() => {
-        expect(seen).toEqual([null, "company-1"]);
-        expect(localStorage.getItem("paperclip.selectedCompanyId")).toBe("company-1");
-      });
-    });
+  it("uses selectable sidebar companies before archived companies", () => {
+    expect(resolveBootstrapCompanySelection({
+      companies: [archivedCompany, activeCompany],
+      sidebarCompanies: [activeCompany],
+      selectedCompanyId: null,
+      storedCompanyId: "archived-company",
+    })).toBe("company-1");
   });
 });

ui/src/context/CompanyContext.tsx

@@ -35,6 +35,26 @@ const STORAGE_KEY = "paperclip.selectedCompanyId";
 
 const CompanyContext = createContext<CompanyContextValue | null>(null);
 
+export function resolveBootstrapCompanySelection(input: {
+  companies: Array<Pick<Company, "id">>;
+  sidebarCompanies: Array<Pick<Company, "id">>;
+  selectedCompanyId: string | null;
+  storedCompanyId: string | null;
+}) {
+  if (input.companies.length === 0) return null;
+
+  const selectableCompanies = input.sidebarCompanies.length > 0
+    ? input.sidebarCompanies
+    : input.companies;
+  if (input.selectedCompanyId && selectableCompanies.some((company) => company.id === input.selectedCompanyId)) {
+    return input.selectedCompanyId;
+  }
+  if (input.storedCompanyId && selectableCompanies.some((company) => company.id === input.storedCompanyId)) {
+    return input.storedCompanyId;
+  }
+  return selectableCompanies[0]?.id ?? null;
+}
+
 export function CompanyProvider({ children }: { children: ReactNode }) {
   const queryClient = useQueryClient();
   const [selectionSource, setSelectionSource] = useState<CompanySelectionSource>("bootstrap");
@@ -70,13 +90,13 @@ export function CompanyProvider({ children }: { children: ReactNode }) {
       return;
     }
 
-    const selectableCompanies = sidebarCompanies.length > 0 ? sidebarCompanies : companies;
-    if (selectedCompanyId && selectableCompanies.some((c) => c.id === selectedCompanyId)) return;
-
-    const stored = localStorage.getItem(STORAGE_KEY);
-    const next = stored && selectableCompanies.some((c) => c.id === stored)
-      ? stored
-      : selectableCompanies[0]!.id;
+    const next = resolveBootstrapCompanySelection({
+      companies,
+      sidebarCompanies,
+      selectedCompanyId,
+      storedCompanyId: localStorage.getItem(STORAGE_KEY),
+    });
+    if (next === null || next === selectedCompanyId) return;
     setSelectedCompanyIdState(next);
     setSelectionSource("bootstrap");
     localStorage.setItem(STORAGE_KEY, next);