64 lines
2.2 KiB
JavaScript
64 lines
2.2 KiB
JavaScript
import { describe, expect, it } from "vitest";
|
|
|
|
import { NODE_CLASSES } from "../constants.js";
|
|
import { markSkipDispatch, selectLeastBusyWorker } from "../executionDecisionUtils.js";
|
|
import { buildWorkerWebSocketUrl } from "../urlUtils.js";
|
|
|
|
|
|
describe("execution decision helpers", () => {
|
|
it("markSkipDispatch sets skip_dispatch=true on DistributedQueue nodes only", () => {
|
|
const prompt = {
|
|
"1": {
|
|
class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
|
|
inputs: { seed: 1 },
|
|
},
|
|
"2": {
|
|
class_type: "KSampler",
|
|
inputs: { seed: 2 },
|
|
},
|
|
"3": {
|
|
class_type: NODE_CLASSES.DISTRIBUTED_QUEUE,
|
|
inputs: {},
|
|
},
|
|
};
|
|
|
|
markSkipDispatch(prompt);
|
|
|
|
expect(prompt["1"].inputs.skip_dispatch).toBe(true);
|
|
expect(prompt["3"].inputs.skip_dispatch).toBe(true);
|
|
expect(prompt["2"].inputs.skip_dispatch).toBeUndefined();
|
|
});
|
|
|
|
it("selectLeastBusyWorker round-robins among idle workers", () => {
|
|
const extension = {};
|
|
const statuses = [
|
|
{ worker: { id: "w1" }, queueRemaining: 0 },
|
|
{ worker: { id: "w2" }, queueRemaining: 0 },
|
|
];
|
|
|
|
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
|
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
|
|
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[0]);
|
|
});
|
|
|
|
it("selectLeastBusyWorker chooses shortest queue when no idle workers", () => {
|
|
const extension = {};
|
|
const statuses = [
|
|
{ worker: { id: "w1" }, queueRemaining: 5 },
|
|
{ worker: { id: "w2" }, queueRemaining: 2 },
|
|
{ worker: { id: "w3" }, queueRemaining: 7 },
|
|
];
|
|
|
|
expect(selectLeastBusyWorker(extension, statuses)).toBe(statuses[1]);
|
|
});
|
|
|
|
it("buildWorkerWebSocketUrl converts http/https to ws/wss", () => {
|
|
expect(buildWorkerWebSocketUrl("http://worker.local:8188")).toBe(
|
|
"ws://worker.local:8188/distributed/worker_ws"
|
|
);
|
|
expect(buildWorkerWebSocketUrl("https://worker.example.com")).toBe(
|
|
"wss://worker.example.com/distributed/worker_ws"
|
|
);
|
|
});
|
|
});
|