Jason WoltjeandClaude Opus 4.5
6a4f58dc1c
- Use SCAN with cursor for non-blocking iteration - Prevents Redis DoS under high key counts - Same API, safer implementation Refs #337 Co-Authored-By: Claude Opus 4.5 <[email protected]>
577 lines
18 KiB
TypeScript
577 lines
18 KiB
TypeScript
import { describe, it, expect, beforeEach, vi, afterEach } from "vitest";
|
|
import { ValkeyClient } from "./valkey.client";
|
|
import type { TaskState, AgentState, OrchestratorEvent } from "./types";
|
|
|
|
// Create a shared mock instance that will be used across all tests
|
|
const mockRedisInstance = {
|
|
get: vi.fn(),
|
|
set: vi.fn(),
|
|
del: vi.fn(),
|
|
publish: vi.fn(),
|
|
subscribe: vi.fn(),
|
|
on: vi.fn(),
|
|
quit: vi.fn(),
|
|
duplicate: vi.fn(),
|
|
scan: vi.fn(),
|
|
};
|
|
|
|
// Mock ioredis
|
|
vi.mock("ioredis", () => {
|
|
return {
|
|
default: class {
|
|
constructor() {
|
|
return mockRedisInstance;
|
|
}
|
|
},
|
|
};
|
|
});
|
|
|
|
describe("ValkeyClient", () => {
|
|
let client: ValkeyClient;
|
|
let mockRedis: typeof mockRedisInstance;
|
|
|
|
beforeEach(() => {
|
|
// Reset all mocks
|
|
vi.clearAllMocks();
|
|
|
|
// Create client instance
|
|
client = new ValkeyClient({
|
|
host: "localhost",
|
|
port: 6379,
|
|
});
|
|
|
|
// Reference the mock instance
|
|
mockRedis = mockRedisInstance;
|
|
|
|
// Mock duplicate to return another mock client
|
|
mockRedis.duplicate.mockReturnValue(mockRedis);
|
|
});
|
|
|
|
afterEach(() => {
|
|
vi.clearAllMocks();
|
|
});
|
|
|
|
describe("Connection Management", () => {
|
|
it("should disconnect on close", async () => {
|
|
mockRedis.quit.mockResolvedValue("OK");
|
|
|
|
await client.disconnect();
|
|
|
|
expect(mockRedis.quit).toHaveBeenCalled();
|
|
});
|
|
|
|
it("should disconnect subscriber if it exists", async () => {
|
|
mockRedis.quit.mockResolvedValue("OK");
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
|
|
// Create subscriber
|
|
await client.subscribeToEvents(vi.fn());
|
|
|
|
await client.disconnect();
|
|
|
|
// Should call quit twice (main client and subscriber)
|
|
expect(mockRedis.quit).toHaveBeenCalledTimes(2);
|
|
});
|
|
});
|
|
|
|
describe("Task State Management", () => {
|
|
const mockTaskState: TaskState = {
|
|
taskId: "task-123",
|
|
status: "pending",
|
|
context: {
|
|
repository: "https://github.com/example/repo",
|
|
branch: "main",
|
|
workItems: ["item-1"],
|
|
},
|
|
createdAt: "2026-02-02T10:00:00Z",
|
|
updatedAt: "2026-02-02T10:00:00Z",
|
|
};
|
|
|
|
it("should get task state", async () => {
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(mockTaskState));
|
|
|
|
const result = await client.getTaskState("task-123");
|
|
|
|
expect(mockRedis.get).toHaveBeenCalledWith("orchestrator:task:task-123");
|
|
expect(result).toEqual(mockTaskState);
|
|
});
|
|
|
|
it("should return null for non-existent task", async () => {
|
|
mockRedis.get.mockResolvedValue(null);
|
|
|
|
const result = await client.getTaskState("task-999");
|
|
|
|
expect(result).toBeNull();
|
|
});
|
|
|
|
it("should set task state", async () => {
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
await client.setTaskState(mockTaskState);
|
|
|
|
expect(mockRedis.set).toHaveBeenCalledWith(
|
|
"orchestrator:task:task-123",
|
|
JSON.stringify(mockTaskState)
|
|
);
|
|
});
|
|
|
|
it("should delete task state", async () => {
|
|
mockRedis.del.mockResolvedValue(1);
|
|
|
|
await client.deleteTaskState("task-123");
|
|
|
|
expect(mockRedis.del).toHaveBeenCalledWith("orchestrator:task:task-123");
|
|
});
|
|
|
|
it("should update task status", async () => {
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(mockTaskState));
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
const result = await client.updateTaskStatus("task-123", "assigned", "agent-456");
|
|
|
|
expect(mockRedis.get).toHaveBeenCalledWith("orchestrator:task:task-123");
|
|
expect(mockRedis.set).toHaveBeenCalled();
|
|
expect(result?.status).toBe("assigned");
|
|
expect(result?.agentId).toBe("agent-456");
|
|
expect(result?.updatedAt).toBeDefined();
|
|
});
|
|
|
|
it("should throw error when updating non-existent task", async () => {
|
|
mockRedis.get.mockResolvedValue(null);
|
|
|
|
await expect(client.updateTaskStatus("task-999", "assigned")).rejects.toThrow(
|
|
"Task task-999 not found"
|
|
);
|
|
});
|
|
|
|
it("should throw error for invalid task status transition", async () => {
|
|
const completedTask = { ...mockTaskState, status: "completed" as const };
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(completedTask));
|
|
|
|
await expect(client.updateTaskStatus("task-123", "assigned")).rejects.toThrow(
|
|
"Invalid task state transition from completed to assigned"
|
|
);
|
|
});
|
|
|
|
it("should list all task states using SCAN", async () => {
|
|
// SCAN returns [cursor, keys] - cursor "0" means complete
|
|
mockRedis.scan.mockResolvedValue([
|
|
"0",
|
|
["orchestrator:task:task-1", "orchestrator:task:task-2"],
|
|
]);
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ ...mockTaskState, taskId: "task-1" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ ...mockTaskState, taskId: "task-2" }));
|
|
|
|
const result = await client.listTasks();
|
|
|
|
expect(mockRedis.scan).toHaveBeenCalledWith(
|
|
"0",
|
|
"MATCH",
|
|
"orchestrator:task:*",
|
|
"COUNT",
|
|
100
|
|
);
|
|
expect(result).toHaveLength(2);
|
|
expect(result[0].taskId).toBe("task-1");
|
|
expect(result[1].taskId).toBe("task-2");
|
|
});
|
|
});
|
|
|
|
describe("Agent State Management", () => {
|
|
const mockAgentState: AgentState = {
|
|
agentId: "agent-456",
|
|
status: "spawning",
|
|
taskId: "task-123",
|
|
};
|
|
|
|
it("should get agent state", async () => {
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(mockAgentState));
|
|
|
|
const result = await client.getAgentState("agent-456");
|
|
|
|
expect(mockRedis.get).toHaveBeenCalledWith("orchestrator:agent:agent-456");
|
|
expect(result).toEqual(mockAgentState);
|
|
});
|
|
|
|
it("should return null for non-existent agent", async () => {
|
|
mockRedis.get.mockResolvedValue(null);
|
|
|
|
const result = await client.getAgentState("agent-999");
|
|
|
|
expect(result).toBeNull();
|
|
});
|
|
|
|
it("should set agent state", async () => {
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
await client.setAgentState(mockAgentState);
|
|
|
|
expect(mockRedis.set).toHaveBeenCalledWith(
|
|
"orchestrator:agent:agent-456",
|
|
JSON.stringify(mockAgentState)
|
|
);
|
|
});
|
|
|
|
it("should delete agent state", async () => {
|
|
mockRedis.del.mockResolvedValue(1);
|
|
|
|
await client.deleteAgentState("agent-456");
|
|
|
|
expect(mockRedis.del).toHaveBeenCalledWith("orchestrator:agent:agent-456");
|
|
});
|
|
|
|
it("should update agent status", async () => {
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(mockAgentState));
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
const result = await client.updateAgentStatus("agent-456", "running");
|
|
|
|
expect(mockRedis.get).toHaveBeenCalledWith("orchestrator:agent:agent-456");
|
|
expect(mockRedis.set).toHaveBeenCalled();
|
|
expect(result?.status).toBe("running");
|
|
expect(result?.startedAt).toBeDefined();
|
|
});
|
|
|
|
it("should set completedAt when status is completed", async () => {
|
|
const runningAgent = { ...mockAgentState, status: "running" as const };
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(runningAgent));
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
const result = await client.updateAgentStatus("agent-456", "completed");
|
|
|
|
expect(result?.status).toBe("completed");
|
|
expect(result?.completedAt).toBeDefined();
|
|
});
|
|
|
|
it("should throw error when updating non-existent agent", async () => {
|
|
mockRedis.get.mockResolvedValue(null);
|
|
|
|
await expect(client.updateAgentStatus("agent-999", "running")).rejects.toThrow(
|
|
"Agent agent-999 not found"
|
|
);
|
|
});
|
|
|
|
it("should throw error for invalid agent status transition", async () => {
|
|
const completedAgent = { ...mockAgentState, status: "completed" as const };
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(completedAgent));
|
|
|
|
await expect(client.updateAgentStatus("agent-456", "running")).rejects.toThrow(
|
|
"Invalid agent state transition from completed to running"
|
|
);
|
|
});
|
|
|
|
it("should list all agent states using SCAN", async () => {
|
|
// SCAN returns [cursor, keys] - cursor "0" means complete
|
|
mockRedis.scan.mockResolvedValue([
|
|
"0",
|
|
["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"],
|
|
]);
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ ...mockAgentState, agentId: "agent-1" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ ...mockAgentState, agentId: "agent-2" }));
|
|
|
|
const result = await client.listAgents();
|
|
|
|
expect(mockRedis.scan).toHaveBeenCalledWith(
|
|
"0",
|
|
"MATCH",
|
|
"orchestrator:agent:*",
|
|
"COUNT",
|
|
100
|
|
);
|
|
expect(result).toHaveLength(2);
|
|
expect(result[0].agentId).toBe("agent-1");
|
|
expect(result[1].agentId).toBe("agent-2");
|
|
});
|
|
});
|
|
|
|
describe("Event Pub/Sub", () => {
|
|
const mockEvent: OrchestratorEvent = {
|
|
type: "agent.spawned",
|
|
agentId: "agent-456",
|
|
taskId: "task-123",
|
|
timestamp: "2026-02-02T10:00:00Z",
|
|
};
|
|
|
|
it("should publish events", async () => {
|
|
mockRedis.publish.mockResolvedValue(1);
|
|
|
|
await client.publishEvent(mockEvent);
|
|
|
|
expect(mockRedis.publish).toHaveBeenCalledWith(
|
|
"orchestrator:events",
|
|
JSON.stringify(mockEvent)
|
|
);
|
|
});
|
|
|
|
it("should subscribe to events", async () => {
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
|
|
const handler = vi.fn();
|
|
await client.subscribeToEvents(handler);
|
|
|
|
expect(mockRedis.duplicate).toHaveBeenCalled();
|
|
expect(mockRedis.subscribe).toHaveBeenCalledWith("orchestrator:events");
|
|
});
|
|
|
|
it("should call handler when event is received", async () => {
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
let messageHandler: ((channel: string, message: string) => void) | undefined;
|
|
|
|
mockRedis.on.mockImplementation(
|
|
(event: string, handler: (channel: string, message: string) => void) => {
|
|
if (event === "message") {
|
|
messageHandler = handler;
|
|
}
|
|
return mockRedis;
|
|
}
|
|
);
|
|
|
|
const handler = vi.fn();
|
|
await client.subscribeToEvents(handler);
|
|
|
|
// Simulate receiving a message
|
|
if (messageHandler) {
|
|
messageHandler("orchestrator:events", JSON.stringify(mockEvent));
|
|
}
|
|
|
|
expect(handler).toHaveBeenCalledWith(mockEvent);
|
|
});
|
|
|
|
it("should handle invalid JSON in events gracefully with logger", async () => {
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
let messageHandler: ((channel: string, message: string) => void) | undefined;
|
|
|
|
mockRedis.on.mockImplementation(
|
|
(event: string, handler: (channel: string, message: string) => void) => {
|
|
if (event === "message") {
|
|
messageHandler = handler;
|
|
}
|
|
return mockRedis;
|
|
}
|
|
);
|
|
|
|
const handler = vi.fn();
|
|
const loggerError = vi.fn();
|
|
|
|
// Create client with logger
|
|
const clientWithLogger = new ValkeyClient({
|
|
host: "localhost",
|
|
port: 6379,
|
|
logger: { error: loggerError },
|
|
});
|
|
|
|
// Mock duplicate for new client
|
|
mockRedis.duplicate.mockReturnValue(mockRedis);
|
|
|
|
await clientWithLogger.subscribeToEvents(handler);
|
|
|
|
// Simulate receiving invalid JSON
|
|
if (messageHandler) {
|
|
messageHandler("orchestrator:events", "invalid json");
|
|
}
|
|
|
|
expect(handler).not.toHaveBeenCalled();
|
|
expect(loggerError).toHaveBeenCalled();
|
|
expect(loggerError).toHaveBeenCalledWith(
|
|
expect.stringContaining("Failed to parse event from channel orchestrator:events"),
|
|
expect.any(Error)
|
|
);
|
|
});
|
|
|
|
it("should invoke error handler when provided", async () => {
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
let messageHandler: ((channel: string, message: string) => void) | undefined;
|
|
|
|
mockRedis.on.mockImplementation(
|
|
(event: string, handler: (channel: string, message: string) => void) => {
|
|
if (event === "message") {
|
|
messageHandler = handler;
|
|
}
|
|
return mockRedis;
|
|
}
|
|
);
|
|
|
|
const handler = vi.fn();
|
|
const errorHandler = vi.fn();
|
|
|
|
await client.subscribeToEvents(handler, errorHandler);
|
|
|
|
// Simulate receiving invalid JSON
|
|
if (messageHandler) {
|
|
messageHandler("orchestrator:events", "invalid json");
|
|
}
|
|
|
|
expect(handler).not.toHaveBeenCalled();
|
|
expect(errorHandler).toHaveBeenCalledWith(
|
|
expect.any(Error),
|
|
"invalid json",
|
|
"orchestrator:events"
|
|
);
|
|
});
|
|
|
|
it("should handle errors without logger or error handler", async () => {
|
|
mockRedis.subscribe.mockResolvedValue(1);
|
|
let messageHandler: ((channel: string, message: string) => void) | undefined;
|
|
|
|
mockRedis.on.mockImplementation(
|
|
(event: string, handler: (channel: string, message: string) => void) => {
|
|
if (event === "message") {
|
|
messageHandler = handler;
|
|
}
|
|
return mockRedis;
|
|
}
|
|
);
|
|
|
|
const handler = vi.fn();
|
|
|
|
await client.subscribeToEvents(handler);
|
|
|
|
// Should not throw when neither logger nor error handler is provided
|
|
expect(() => {
|
|
if (messageHandler) {
|
|
messageHandler("orchestrator:events", "invalid json");
|
|
}
|
|
}).not.toThrow();
|
|
|
|
expect(handler).not.toHaveBeenCalled();
|
|
});
|
|
});
|
|
|
|
describe("Edge Cases", () => {
|
|
it("should handle task updates with error parameter", async () => {
|
|
const taskState: TaskState = {
|
|
taskId: "task-123",
|
|
status: "pending",
|
|
context: {
|
|
repository: "https://github.com/example/repo",
|
|
branch: "main",
|
|
workItems: ["item-1"],
|
|
},
|
|
createdAt: "2026-02-02T10:00:00Z",
|
|
updatedAt: "2026-02-02T10:00:00Z",
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(taskState));
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
const result = await client.updateTaskStatus("task-123", "failed", undefined, "Test error");
|
|
|
|
expect(result.status).toBe("failed");
|
|
expect(result.metadata?.error).toBe("Test error");
|
|
});
|
|
|
|
it("should handle agent updates with error parameter", async () => {
|
|
const agentState: AgentState = {
|
|
agentId: "agent-456",
|
|
status: "running",
|
|
taskId: "task-123",
|
|
};
|
|
|
|
mockRedis.get.mockResolvedValue(JSON.stringify(agentState));
|
|
mockRedis.set.mockResolvedValue("OK");
|
|
|
|
const result = await client.updateAgentStatus("agent-456", "failed", "Test error");
|
|
|
|
expect(result.status).toBe("failed");
|
|
expect(result.error).toBe("Test error");
|
|
});
|
|
|
|
it("should filter out null values in listTasks", async () => {
|
|
mockRedis.scan.mockResolvedValue([
|
|
"0",
|
|
["orchestrator:task:task-1", "orchestrator:task:task-2"],
|
|
]);
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-1", status: "pending" }))
|
|
.mockResolvedValueOnce(null); // Simulate deleted task
|
|
|
|
const result = await client.listTasks();
|
|
|
|
expect(result).toHaveLength(1);
|
|
expect(result[0].taskId).toBe("task-1");
|
|
});
|
|
|
|
it("should filter out null values in listAgents", async () => {
|
|
mockRedis.scan.mockResolvedValue([
|
|
"0",
|
|
["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"],
|
|
]);
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-1", status: "running" }))
|
|
.mockResolvedValueOnce(null); // Simulate deleted agent
|
|
|
|
const result = await client.listAgents();
|
|
|
|
expect(result).toHaveLength(1);
|
|
expect(result[0].agentId).toBe("agent-1");
|
|
});
|
|
});
|
|
|
|
describe("SCAN-based iteration (large key sets)", () => {
|
|
it("should handle multiple SCAN iterations for tasks", async () => {
|
|
// Simulate SCAN returning multiple batches with cursor pagination
|
|
mockRedis.scan
|
|
.mockResolvedValueOnce(["42", ["orchestrator:task:task-1", "orchestrator:task:task-2"]]) // First batch, cursor 42
|
|
.mockResolvedValueOnce(["0", ["orchestrator:task:task-3"]]); // Second batch, cursor 0 = done
|
|
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-1", status: "pending" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-2", status: "pending" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ taskId: "task-3", status: "pending" }));
|
|
|
|
const result = await client.listTasks();
|
|
|
|
expect(mockRedis.scan).toHaveBeenCalledTimes(2);
|
|
expect(mockRedis.scan).toHaveBeenNthCalledWith(
|
|
1,
|
|
"0",
|
|
"MATCH",
|
|
"orchestrator:task:*",
|
|
"COUNT",
|
|
100
|
|
);
|
|
expect(mockRedis.scan).toHaveBeenNthCalledWith(
|
|
2,
|
|
"42",
|
|
"MATCH",
|
|
"orchestrator:task:*",
|
|
"COUNT",
|
|
100
|
|
);
|
|
expect(result).toHaveLength(3);
|
|
expect(result.map((t) => t.taskId)).toEqual(["task-1", "task-2", "task-3"]);
|
|
});
|
|
|
|
it("should handle multiple SCAN iterations for agents", async () => {
|
|
// Simulate SCAN returning multiple batches with cursor pagination
|
|
mockRedis.scan
|
|
.mockResolvedValueOnce(["99", ["orchestrator:agent:agent-1", "orchestrator:agent:agent-2"]]) // First batch
|
|
.mockResolvedValueOnce(["50", ["orchestrator:agent:agent-3"]]) // Second batch
|
|
.mockResolvedValueOnce(["0", ["orchestrator:agent:agent-4"]]); // Third batch, done
|
|
|
|
mockRedis.get
|
|
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-1", status: "running" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-2", status: "running" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-3", status: "running" }))
|
|
.mockResolvedValueOnce(JSON.stringify({ agentId: "agent-4", status: "running" }));
|
|
|
|
const result = await client.listAgents();
|
|
|
|
expect(mockRedis.scan).toHaveBeenCalledTimes(3);
|
|
expect(result).toHaveLength(4);
|
|
expect(result.map((a) => a.agentId)).toEqual(["agent-1", "agent-2", "agent-3", "agent-4"]);
|
|
});
|
|
|
|
it("should handle empty result from SCAN", async () => {
|
|
mockRedis.scan.mockResolvedValue(["0", []]);
|
|
|
|
const result = await client.listTasks();
|
|
|
|
expect(mockRedis.scan).toHaveBeenCalledTimes(1);
|
|
expect(result).toHaveLength(0);
|
|
});
|
|
});
|
|
});
|