import { describe, it, expect, vi, afterEach } from "vitest"; import { update_object, walk_and_store_blobs, skip_queue, post_message, handle_file, handle_payload } from "../helpers/data"; import { config_response, endpoint_info } from "./test_data"; import { BlobRef, Command } from "../types"; import { FileData } from "../upload"; const IS_NODE = process.env.TEST_MODE === "node"; describe("walk_and_store_blobs", () => { it("should convert a Buffer to a Blob", async () => { const buffer = Buffer.from("test data"); const parts = await walk_and_store_blobs(buffer, "text"); expect(parts).toHaveLength(1); expect(parts[0].blob).toBeInstanceOf(Blob); }); it("should return a Blob when passed a Blob", async () => { const blob = new Blob(["test data"]); const parts = await walk_and_store_blobs( blob, undefined, [], true, endpoint_info ); expect(parts[0].blob).toBeInstanceOf(Blob); }); it("should handle arrays", async () => { const image = new Blob([]); const parts = await walk_and_store_blobs([image]); expect(parts).toHaveLength(1); expect(parts[0].blob).toBeInstanceOf(Blob); expect(parts[0].path).toEqual(["0"]); }); it("should handle deep structures", async () => { const image = new Blob([]); const parts = await walk_and_store_blobs({ a: { b: { data: { image } } } }); expect(parts).toHaveLength(1); expect(parts[0].blob).toBeInstanceOf(Blob); expect(parts[0].path).toEqual(["a", "b", "data", "image"]); }); it("should handle deep structures with arrays", async () => { const image = new Blob([]); const parts = await walk_and_store_blobs({ a: [ { b: [ { data: [ { image } ] } ] } ] }); expect(parts[0].blob).toBeInstanceOf(Blob); }); it("should handle deep structures with arrays (with equality check)", async () => { const image = new Blob([]); const obj = { a: [ { b: [ { data: [[image], image, [image, [image]]] } ] } ] }; const parts = await walk_and_store_blobs(obj); async function map_path(obj: Record, parts: BlobRef[]) { const { path, blob } = parts[parts.length - 1]; let ref = obj; path.forEach((p) => (ref = ref[p])); // since ref is a Blob and blob is a Blob, we deep equal check the two buffers instead if (ref instanceof Blob && blob instanceof Blob) { const refBuffer = Buffer.from(await ref.arrayBuffer()); const blobBuffer = Buffer.from(await blob.arrayBuffer()); return refBuffer.equals(blobBuffer); } return ref === blob; } expect(parts[0].blob).toBeInstanceOf(Blob); expect(map_path(obj, parts)).toBeTruthy(); }); it("should handle buffer instances and return a BlobRef", async () => { const buffer = Buffer.from("test"); const parts = await walk_and_store_blobs(buffer, undefined, ["blob"]); expect(parts).toHaveLength(1); expect(parts[0].blob).toBeInstanceOf(Blob); expect(parts[0].path).toEqual(["blob"]); }); it("should handle buffer instances with a path and return a BlobRef with the path", async () => { const buffer = Buffer.from("test data"); const parts = await walk_and_store_blobs(buffer); expect(parts).toHaveLength(1); expect(parts[0].path).toEqual([]); expect(parts[0].blob).toBeInstanceOf(Blob); }); it("should convert an object with deep structures to BlobRefs", async () => { const param = { a: { b: { data: { image: Buffer.from("test image") } } } }; const parts = await walk_and_store_blobs(param); expect(parts).toHaveLength(1); expect(parts[0].path).toEqual(["a", "b", "data", "image"]); expect(parts[0].blob).toBeInstanceOf(Blob); }); }); describe("update_object", () => { it("should update the value of a nested property", () => { const obj = { a: { b: { c: "old value" } } }; const stack = ["a", "b", "c"]; const new_val = "new value"; update_object(obj, new_val, stack); expect(obj.a.b.c).toBe(new_val); }); it("should throw an error for invalid key type", () => { const obj = { a: { b: { c: "value" } } }; const stack = ["a", "b", true]; const newValue = "new value"; expect(() => { // @ts-ignore update_object(obj, newValue, stack); }).toThrowError("Invalid key type"); }); }); describe("skip_queue", () => { const id = 0; const config = config_response; it("should not skip queue when global and dependency queue is enabled", () => { config.enable_queue = true; config.dependencies.find((dep) => dep.id === id)!.queue = true; const result = skip_queue(id, config_response); expect(result).toBe(false); }); it("should not skip queue when global queue is disabled and dependency queue is enabled", () => { config.enable_queue = false; config.dependencies.find((dep) => dep.id === id)!.queue = true; const result = skip_queue(id, config_response); expect(result).toBe(false); }); it("should should skip queue when global queue and dependency queue is disabled", () => { config.enable_queue = false; config.dependencies.find((dep) => dep.id === id)!.queue = false; const result = skip_queue(id, config_response); expect(result).toBe(true); }); it("should should skip queue when global queue is enabled and dependency queue is disabled", () => { config.enable_queue = true; config.dependencies.find((dep) => dep.id === id)!.queue = false; const result = skip_queue(id, config_response); expect(result).toBe(true); }); }); describe("post_message", () => { afterEach(() => { vi.restoreAllMocks(); }); it("should send a message to the parent window and resolve with received data", async () => { const test_data = { key: "value" }; const test_origin = "https://huggingface.co"; const post_message_mock = vi.fn(); global.window = { // @ts-ignore parent: { postMessage: post_message_mock } }; const message_channel_mock = { port1: { onmessage: (handler) => { onmessage = handler; }, close: vi.fn() }, port2: {} }; vi.stubGlobal("MessageChannel", function () { this.port1 = message_channel_mock.port1; this.port2 = message_channel_mock.port2; return this; }); const promise = post_message(test_data, test_origin); if (message_channel_mock.port1.onmessage) { message_channel_mock.port1.onmessage({ data: test_data }); } await expect(promise).resolves.toEqual(test_data); expect(post_message_mock).toHaveBeenCalledWith(test_data, test_origin, [ message_channel_mock.port2 ]); }); }); describe("handle_file", () => { it("should handle a Blob object and return the blob", () => { const blob = new Blob(["test data"], { type: "image/png" }); const result = handle_file(blob) as FileData; expect(result).toBe(blob); }); it("should handle a Buffer object and return it as a blob", () => { const buffer = Buffer.from("test data"); const result = handle_file(buffer) as FileData; expect(result).toBeInstanceOf(Blob); }); it("should handle a local file path and return a Command object", () => { const file_path = "./owl.png"; const result = handle_file(file_path) as Command; expect(result).toBeInstanceOf(Command); expect(result).toEqual({ type: "command", command: "upload_file", meta: { path: "./owl.png", name: "./owl.png", orig_path: "./owl.png" }, fileData: undefined }); }); it("should handle a File object and return it as FileData", () => { if (IS_NODE) { return; } const file = new File(["test image"], "test.png", { type: "image/png" }); const result = handle_file(file) as FileData; expect(result).toBeInstanceOf(Blob); }); it("should throw an error for invalid input", () => { const invalid_input = 123; expect(() => { // @ts-ignore handle_file(invalid_input); }).toThrowError( "Invalid input: must be a URL, File, Blob, or Buffer object." ); }); }); describe("handle_payload", () => { it("should return an input payload with null in place of `state` when with_null_state is true", () => { const resolved_payload = [2]; const dependency = { inputs: [1, 2] }; const components = [ { id: 1, type: "number" }, { id: 2, type: "state" } ]; const with_null_state = true; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "input", with_null_state ); expect(result).toEqual([2, null]); }); it("should return an input payload with null in place of two `state` components when with_null_state is true", () => { const resolved_payload = ["hello", "goodbye"]; const dependency = { inputs: [1, 2, 3, 4] }; const components = [ { id: 1, type: "textbox" }, { id: 2, type: "state" }, { id: 3, type: "textbox" }, { id: 4, type: "state" } ]; const with_null_state = true; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "input", with_null_state ); expect(result).toEqual(["hello", null, "goodbye", null]); }); it("should return an output payload without the state component value when with_null_state is false", () => { const resolved_payload = ["hello", null]; const dependency = { outputs: [2, 3] }; const components = [ { id: 2, type: "textbox" }, { id: 3, type: "state" } ]; const with_null_state = false; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "output", with_null_state ); expect(result).toEqual(["hello"]); }); it("should return an ouput payload without the two state component values when with_null_state is false", () => { const resolved_payload = ["hello", null, "world", null]; const dependency = { outputs: [2, 3, 4, 5] }; const components = [ { id: 2, type: "textbox" }, { id: 3, type: "state" }, { id: 4, type: "textbox" }, { id: 5, type: "state" } ]; const with_null_state = false; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "output", with_null_state ); expect(result).toEqual(["hello", "world"]); }); it("should return an ouput payload with the two state component values when with_null_state is true", () => { const resolved_payload = ["hello", null, "world", null]; const dependency = { outputs: [2, 3, 4, 5] }; const components = [ { id: 2, type: "textbox" }, { id: 3, type: "state" }, { id: 4, type: "textbox" }, { id: 5, type: "state" } ]; const with_null_state = true; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "output", with_null_state ); expect(result).toEqual(["hello", null, "world", null]); }); it("should return the same payload where no state components are defined", () => { const resolved_payload = ["hello", "world"]; const dependency = { inputs: [2, 3] }; const components = [ { id: 2, type: "textbox" }, { id: 3, type: "textbox" } ]; const with_null_state = true; const result = handle_payload( resolved_payload, // @ts-ignore dependency, components, "input", with_null_state ); expect(result).toEqual(["hello", "world"]); }); });