import { QueryClient } from "react-query";

import { describe, expect, it } from "vitest";

import type { WorkspacesResponse } from "#/api/typesGenerated";

import { MockWorkspace } from "#/testHelpers/entities";

import { getWorkspaceQuotaQueryKey } from "./workspaceQuota";

import {

autoCreateWorkspace,

buildLogsKey,

createWorkspace,

invalidateWorkspaceListQueries,

invalidateWorkspaceMutationQueries,

setOptimisticWorkspaceListBuildStatus,

workspacesKey,

workspacesQueryKeyPrefix,

workspaceUsage,

} from "./workspaces";

const createTestQueryClient = (): QueryClient =>

new QueryClient({

defaultOptions: {

queries: {

retry: false,

gcTime: Number.POSITIVE_INFINITY,

refetchOnWindowFocus: false,

networkMode: "offlineFirst",

},

},

});

const workspacesResponse = {

workspaces: [],

count: 0,

} satisfies WorkspacesResponse;

const seedWorkspaceFamilyQueries = (queryClient: QueryClient) => {

const rawListKey = workspacesQueryKeyPrefix;

const defaultListKey = workspacesKey({});

const filteredListKey = workspacesKey({

q: "owner:me organization:default",

limit: 25,

offset: 50,

});

const usageKey = workspaceUsage({

usageApp: "reconnecting-pty",

connectionStatus: "connected",

workspaceId: "workspace-1",

agentId: "agent-1",

}).queryKey;

const buildLogs = buildLogsKey("workspace-1");

const workspacePermissionsKey = [

"workspaces",

"workspace-1",

"permissions",

] as const;

const workspaceAgentCredentialsKey = [

"workspaces",

"workspace-1",

"agents",

"main",

"credentials",

] as const;

const organizationWorkspacePermissionsKey = [

"workspaces",

["organization-1"],

"permissions",

] as const;

queryClient.setQueryData(rawListKey, workspacesResponse);

queryClient.setQueryData(defaultListKey, workspacesResponse);

queryClient.setQueryData(filteredListKey, workspacesResponse);

queryClient.setQueryData(usageKey, { tracked: true });

queryClient.setQueryData(buildLogs, []);

queryClient.setQueryData(workspacePermissionsKey, { read: true });

queryClient.setQueryData(workspaceAgentCredentialsKey, { token: "secret" });

queryClient.setQueryData(organizationWorkspacePermissionsKey, { read: true });

return {

listKeys: [rawListKey, defaultListKey, filteredListKey],

nonListKeys: [

usageKey,

buildLogs,

workspacePermissionsKey,

workspaceAgentCredentialsKey,

organizationWorkspacePermissionsKey,

],

};

};

describe("invalidateWorkspaceListQueries", () => {

it("invalidates workspace list queries without touching side-effecting workspace-family queries", async () => {

const queryClient = createTestQueryClient();

const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient);

await invalidateWorkspaceListQueries(queryClient);

for (const key of listKeys) {

expect(

queryClient.getQueryState(key)?.isInvalidated,

`${JSON.stringify(key)} should be invalidated`,

).toBe(true);

}

for (const key of nonListKeys) {

expect(

queryClient.getQueryState(key)?.isInvalidated,

`${JSON.stringify(key)} should NOT be invalidated`,

).not.toBe(true);

}

});

});

describe("invalidateWorkspaceMutationQueries", () => {

it("uses narrowed list invalidation and keeps workspace usage queries untouched", async () => {

const queryClient = createTestQueryClient();

const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient);

const quotaKey = getWorkspaceQuotaQueryKey("default", "me");

queryClient.setQueryData(quotaKey, { credits_consumed: 1, budget: 10 });

await invalidateWorkspaceMutationQueries(queryClient, {

organizationName: "default",

username: "me",

});

for (const key of listKeys) {

expect(

queryClient.getQueryState(key)?.isInvalidated,

`${JSON.stringify(key)} should be invalidated`,

).toBe(true);

}

expect(queryClient.getQueryState(quotaKey)?.isInvalidated).toBe(true);

for (const key of nonListKeys) {

expect(

queryClient.getQueryState(key)?.isInvalidated,

`${JSON.stringify(key)} should NOT be invalidated`,

).not.toBe(true);

}

});

});

describe("setOptimisticWorkspaceListBuildStatus", () => {

it("patches the matching workspace's build status across list caches and rolls back", () => {

const queryClient = createTestQueryClient();

const listKey = workspacesKey({});

const seeded = {

workspaces: [MockWorkspace],

count: 1,

} satisfies WorkspacesResponse;

queryClient.setQueryData(listKey, seeded);

const rollback = setOptimisticWorkspaceListBuildStatus(

queryClient,

MockWorkspace.id,

"stopping",

"stop",

);

const patched = queryClient.getQueryData<WorkspacesResponse>(listKey);

expect(patched?.workspaces[0].latest_build.status).toBe("stopping");

expect(patched?.workspaces[0].latest_build.transition).toBe("stop");

rollback();

const restored = queryClient.getQueryData<WorkspacesResponse>(listKey);

expect(restored?.workspaces[0].latest_build.status).toBe("running");

});

});

describe("workspace creation mutations", () => {

it("use narrowed list invalidation for manual workspace creation", async () => {

const queryClient = createTestQueryClient();

const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient);

await createWorkspace(queryClient).onSuccess();

for (const key of listKeys) {

expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true);

}

for (const key of nonListKeys) {

expect(queryClient.getQueryState(key)?.isInvalidated).not.toBe(true);

}

});

it("use narrowed list invalidation for auto workspace creation", async () => {

const queryClient = createTestQueryClient();

const { listKeys, nonListKeys } = seedWorkspaceFamilyQueries(queryClient);

await autoCreateWorkspace(queryClient).onSuccess();

for (const key of listKeys) {

expect(queryClient.getQueryState(key)?.isInvalidated).toBe(true);

}

for (const key of nonListKeys) {

expect(queryClient.getQueryState(key)?.isInvalidated).not.toBe(true);

}

});

});