From 1bd1cac23fabefdf9da29c79f183e31221611331 Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Mon, 20 Apr 2026 21:35:27 +0100 Subject: [PATCH] test: share comfy provider fixtures --- .../comfy/image-generation-provider.test.ts | 83 +++------------- extensions/comfy/test-helpers.ts | 99 +++++++++++++++++++ .../comfy/video-generation-provider.test.ts | 83 +++------------- 3 files changed, 131 insertions(+), 134 deletions(-) create mode 100644 extensions/comfy/test-helpers.ts diff --git a/extensions/comfy/image-generation-provider.test.ts b/extensions/comfy/image-generation-provider.test.ts index 3fec5f228d0..f1fc4b9a8d3 100644 --- a/extensions/comfy/image-generation-provider.test.ts +++ b/extensions/comfy/image-generation-provider.test.ts @@ -1,29 +1,21 @@ -import type { OpenClawConfig } from "openclaw/plugin-sdk/config-runtime"; -import * as providerAuth from "openclaw/plugin-sdk/provider-auth-runtime"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { _setComfyFetchGuardForTesting, buildComfyImageGenerationProvider, } from "./image-generation-provider.js"; +import { + buildComfyConfig, + mockComfyCloudJobResponses, + mockComfyProviderApiKey, + parseComfyJsonBody, +} from "./test-helpers.js"; const { fetchWithSsrFGuardMock } = vi.hoisted(() => ({ fetchWithSsrFGuardMock: vi.fn(), })); function parseJsonBody(call: number): Record { - const request = fetchWithSsrFGuardMock.mock.calls[call - 1]?.[0]; - expect(request?.init?.body).toBeTruthy(); - return JSON.parse(String(request.init.body)) as Record; -} - -function buildComfyConfig(config: Record): OpenClawConfig { - return { - models: { - providers: { - comfy: config, - }, - }, - } as unknown as OpenClawConfig; + return parseComfyJsonBody(fetchWithSsrFGuardMock, call); } describe("comfy image-generation provider", () => { @@ -234,59 +226,16 @@ describe("comfy image-generation provider", () => { }); it("uses cloud endpoints, auth headers, and partner-node extra_data", async () => { - vi.spyOn(providerAuth, "resolveApiKeyForProvider").mockResolvedValue({ - apiKey: "comfy-test-key", - source: "env", - mode: "api-key", - }); + mockComfyProviderApiKey(); _setComfyFetchGuardForTesting(fetchWithSsrFGuardMock); - fetchWithSsrFGuardMock - .mockResolvedValueOnce({ - response: new Response(JSON.stringify({ prompt_id: "cloud-job-1" }), { - status: 200, - headers: { "content-type": "application/json" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(JSON.stringify({ status: "completed" }), { - status: 200, - headers: { "content-type": "application/json" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response( - JSON.stringify({ - "cloud-job-1": { - outputs: { - "9": { - images: [{ filename: "cloud.png", subfolder: "", type: "output" }], - }, - }, - }, - }), - { - status: 200, - headers: { "content-type": "application/json" }, - }, - ), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(null, { - status: 302, - headers: { location: "https://cdn.example.com/cloud.png" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(Buffer.from("cloud-data"), { - status: 200, - headers: { "content-type": "image/png" }, - }), - release: vi.fn(async () => {}), - }); + mockComfyCloudJobResponses(fetchWithSsrFGuardMock, { + body: Buffer.from("cloud-data"), + contentType: "image/png", + filename: "cloud.png", + outputKind: "images", + promptId: "cloud-job-1", + redirectLocation: "https://cdn.example.com/cloud.png", + }); const provider = buildComfyImageGenerationProvider(); const result = await provider.generateImage({ diff --git a/extensions/comfy/test-helpers.ts b/extensions/comfy/test-helpers.ts new file mode 100644 index 00000000000..fe19507827d --- /dev/null +++ b/extensions/comfy/test-helpers.ts @@ -0,0 +1,99 @@ +import type { OpenClawConfig } from "openclaw/plugin-sdk/config-runtime"; +import * as providerAuth from "openclaw/plugin-sdk/provider-auth-runtime"; +import { expect, vi } from "vitest"; + +type FetchGuardMock = ReturnType; + +type FetchGuardRequest = { + init?: { + body?: unknown; + }; +}; + +type ComfyCloudJobResponseOptions = { + body: BodyInit; + contentType: string; + filename: string; + outputKind: "gifs" | "images"; + promptId: string; + redirectLocation: string; +}; + +export function buildComfyConfig(config: Record): OpenClawConfig { + return { + models: { + providers: { + comfy: config, + }, + }, + } as unknown as OpenClawConfig; +} + +export function parseComfyJsonBody( + fetchWithSsrFGuardMock: FetchGuardMock, + call: number, +): Record { + const request = fetchWithSsrFGuardMock.mock.calls[call - 1]?.[0] as FetchGuardRequest | undefined; + expect(request?.init?.body).toBeTruthy(); + return JSON.parse(String(request.init.body)) as Record; +} + +export function mockComfyProviderApiKey(apiKey = "comfy-test-key") { + return vi.spyOn(providerAuth, "resolveApiKeyForProvider").mockResolvedValue({ + apiKey, + source: "env", + mode: "api-key", + }); +} + +export function mockComfyCloudJobResponses( + fetchWithSsrFGuardMock: FetchGuardMock, + options: ComfyCloudJobResponseOptions, +) { + fetchWithSsrFGuardMock + .mockResolvedValueOnce(fetchGuardJson({ prompt_id: options.promptId })) + .mockResolvedValueOnce(fetchGuardJson({ status: "completed" })) + .mockResolvedValueOnce( + fetchGuardJson({ + [options.promptId]: { + outputs: { + "9": { + [options.outputKind]: [{ filename: options.filename, subfolder: "", type: "output" }], + }, + }, + }, + }), + ) + .mockResolvedValueOnce( + fetchGuardResponse( + new Response(null, { + status: 302, + headers: { location: options.redirectLocation }, + }), + ), + ) + .mockResolvedValueOnce( + fetchGuardResponse( + new Response(options.body, { + status: 200, + headers: { "content-type": options.contentType }, + }), + ), + ); +} + +function fetchGuardJson(body: unknown) { + return fetchGuardResponse( + new Response(JSON.stringify(body), { + status: 200, + headers: { "content-type": "application/json" }, + }), + ); +} + +function fetchGuardResponse(response: Response) { + return { + response, + release: vi.fn(async () => {}), + }; +} diff --git a/extensions/comfy/video-generation-provider.test.ts b/extensions/comfy/video-generation-provider.test.ts index 1a5cefc6cc4..550d8a294be 100644 --- a/extensions/comfy/video-generation-provider.test.ts +++ b/extensions/comfy/video-generation-provider.test.ts @@ -1,7 +1,11 @@ -import type { OpenClawConfig } from "openclaw/plugin-sdk/config-runtime"; -import * as providerAuth from "openclaw/plugin-sdk/provider-auth-runtime"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { expectExplicitVideoGenerationCapabilities } from "../../test/helpers/media-generation/provider-capability-assertions.js"; +import { + buildComfyConfig, + mockComfyCloudJobResponses, + mockComfyProviderApiKey, + parseComfyJsonBody, +} from "./test-helpers.js"; import { _setComfyFetchGuardForTesting, buildComfyVideoGenerationProvider, @@ -12,19 +16,7 @@ const { fetchWithSsrFGuardMock } = vi.hoisted(() => ({ })); function parseJsonBody(call: number): Record { - const request = fetchWithSsrFGuardMock.mock.calls[call - 1]?.[0]; - expect(request?.init?.body).toBeTruthy(); - return JSON.parse(String(request.init.body)) as Record; -} - -function buildComfyConfig(config: Record): OpenClawConfig { - return { - models: { - providers: { - comfy: config, - }, - }, - } as unknown as OpenClawConfig; + return parseComfyJsonBody(fetchWithSsrFGuardMock, call); } describe("comfy video-generation provider", () => { @@ -158,59 +150,16 @@ describe("comfy video-generation provider", () => { }); it("uses cloud endpoints for video workflows", async () => { - vi.spyOn(providerAuth, "resolveApiKeyForProvider").mockResolvedValue({ - apiKey: "comfy-test-key", - source: "env", - mode: "api-key", - }); + mockComfyProviderApiKey(); _setComfyFetchGuardForTesting(fetchWithSsrFGuardMock); - fetchWithSsrFGuardMock - .mockResolvedValueOnce({ - response: new Response(JSON.stringify({ prompt_id: "cloud-video-1" }), { - status: 200, - headers: { "content-type": "application/json" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(JSON.stringify({ status: "completed" }), { - status: 200, - headers: { "content-type": "application/json" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response( - JSON.stringify({ - "cloud-video-1": { - outputs: { - "9": { - gifs: [{ filename: "cloud.mp4", subfolder: "", type: "output" }], - }, - }, - }, - }), - { - status: 200, - headers: { "content-type": "application/json" }, - }, - ), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(null, { - status: 302, - headers: { location: "https://cdn.example.com/cloud.mp4" }, - }), - release: vi.fn(async () => {}), - }) - .mockResolvedValueOnce({ - response: new Response(Buffer.from("cloud-video-data"), { - status: 200, - headers: { "content-type": "video/mp4" }, - }), - release: vi.fn(async () => {}), - }); + mockComfyCloudJobResponses(fetchWithSsrFGuardMock, { + body: Buffer.from("cloud-video-data"), + contentType: "video/mp4", + filename: "cloud.mp4", + outputKind: "gifs", + promptId: "cloud-video-1", + redirectLocation: "https://cdn.example.com/cloud.mp4", + }); const provider = buildComfyVideoGenerationProvider(); const result = await provider.generateVideo({