|
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<string, any>, parts: BlobRef[]) {
|
|
const { path, blob } = parts[parts.length - 1];
|
|
let ref = obj;
|
|
path.forEach((p) => (ref = ref[p]));
|
|
|
|
|
|
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(() => {
|
|
|
|
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 = {
|
|
|
|
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(() => {
|
|
|
|
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,
|
|
|
|
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,
|
|
|
|
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,
|
|
|
|
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,
|
|
|
|
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,
|
|
|
|
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,
|
|
|
|
dependency,
|
|
components,
|
|
"input",
|
|
with_null_state
|
|
);
|
|
expect(result).toEqual(["hello", "world"]);
|
|
});
|
|
});
|
|
|