Skip to content

Commit e59e816

Browse files
joshsnyjoshsny
committed
test infrastructure for agent server
1 parent 9c61e31 commit e59e816

9 files changed

Lines changed: 791 additions & 0 deletions

File tree

‎packages/agent/src/test/assertions.ts‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import { expect } from "vitest";
2+
3+
export interface NotificationEntry {
4+
notification?: {
5+
method?: string;
6+
params?: {
7+
sessionId?: string;
8+
update?: {
9+
sessionUpdate?: string;
10+
content?: {
11+
type?: string;
12+
text?: string;
13+
};
14+
};
15+
};
16+
};
17+
}
18+
19+
export interface NotificationMatcher {
20+
method?: string;
21+
text?: string;
22+
sessionId?: string;
23+
sessionUpdate?: string;
24+
}
25+
26+
function entryMatchesNotification(
27+
entry: unknown,
28+
matcher: NotificationMatcher,
29+
): boolean {
30+
const notification = (entry as NotificationEntry).notification;
31+
if (!notification) return false;
32+
33+
if (matcher.method && notification.method !== matcher.method) {
34+
return false;
35+
}
36+
if (
37+
matcher.sessionId &&
38+
notification.params?.sessionId !== matcher.sessionId
39+
) {
40+
return false;
41+
}
42+
if (
43+
matcher.sessionUpdate &&
44+
notification.params?.update?.sessionUpdate !== matcher.sessionUpdate
45+
) {
46+
return false;
47+
}
48+
if (matcher.text) {
49+
const text = notification.params?.update?.content?.text;
50+
if (!text || !text.includes(matcher.text)) {
51+
return false;
52+
}
53+
}
54+
return true;
55+
}
56+
57+
export function findNotification(
58+
appendLogCalls: unknown[][],
59+
matcher: NotificationMatcher,
60+
): NotificationEntry | undefined {
61+
for (const entries of appendLogCalls) {
62+
for (const entry of entries) {
63+
if (entryMatchesNotification(entry, matcher)) {
64+
return entry as NotificationEntry;
65+
}
66+
}
67+
}
68+
return undefined;
69+
}
70+
71+
export function hasNotification(
72+
appendLogCalls: unknown[][],
73+
matcher: NotificationMatcher,
74+
): boolean {
75+
return findNotification(appendLogCalls, matcher) !== undefined;
76+
}
77+
78+
export function expectNotification(
79+
appendLogCalls: unknown[][],
80+
matcher: NotificationMatcher,
81+
): NotificationEntry {
82+
const found = findNotification(appendLogCalls, matcher);
83+
expect(
84+
found,
85+
`Expected notification matching ${JSON.stringify(matcher)}`,
86+
).toBeDefined();
87+
return found!;
88+
}
89+
90+
export function expectNoNotification(
91+
appendLogCalls: unknown[][],
92+
matcher: NotificationMatcher,
93+
): void {
94+
const found = findNotification(appendLogCalls, matcher);
95+
expect(
96+
found,
97+
`Expected no notification matching ${JSON.stringify(matcher)}`,
98+
).toBeUndefined();
99+
}
100+
101+
export function countNotifications(
102+
appendLogCalls: unknown[][],
103+
matcher: NotificationMatcher,
104+
): number {
105+
let count = 0;
106+
for (const entries of appendLogCalls) {
107+
for (const entry of entries) {
108+
if (entryMatchesNotification(entry, matcher)) {
109+
count++;
110+
}
111+
}
112+
}
113+
return count;
114+
}

‎packages/agent/src/test/controllers/sse-controller.ts‎

Lines changed: 107 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
type SseState = "idle" | "streaming" | "closing" | "closed";
2+
3+
export interface EventOptions {
4+
id?: string;
5+
event?: string;
6+
}
7+
8+
export class SseController {
9+
private encoder = new TextEncoder();
10+
private controller: ReadableStreamDefaultController<Uint8Array> | null = null;
11+
private state: SseState = "idle";
12+
13+
createStream(): ReadableStream<Uint8Array> {
14+
return new ReadableStream({
15+
start: (controller) => {
16+
this.controller = controller;
17+
this.state = "streaming";
18+
},
19+
cancel: () => {
20+
this.state = "closed";
21+
},
22+
});
23+
}
24+
25+
sendEvent(data: unknown, options: EventOptions = {}): boolean {
26+
if (
27+
this.state === "closed" ||
28+
this.state === "closing" ||
29+
!this.controller
30+
) {
31+
return false;
32+
}
33+
34+
const lines: string[] = [];
35+
if (options.id) {
36+
lines.push(`id: ${options.id}`);
37+
}
38+
if (options.event) {
39+
lines.push(`event: ${options.event}`);
40+
}
41+
lines.push(`data: ${JSON.stringify(data)}`);
42+
lines.push("");
43+
lines.push("");
44+
45+
try {
46+
this.controller.enqueue(this.encoder.encode(lines.join("\n")));
47+
return true;
48+
} catch {
49+
this.state = "closed";
50+
return false;
51+
}
52+
}
53+
54+
sendRaw(rawData: string): boolean {
55+
if (
56+
this.state === "closed" ||
57+
this.state === "closing" ||
58+
!this.controller
59+
) {
60+
return false;
61+
}
62+
try {
63+
this.controller.enqueue(this.encoder.encode(rawData));
64+
return true;
65+
} catch {
66+
this.state = "closed";
67+
return false;
68+
}
69+
}
70+
71+
sendPartial(partialData: string): boolean {
72+
return this.sendRaw(partialData);
73+
}
74+
75+
error(err: Error): void {
76+
if (this.state !== "streaming" || !this.controller) return;
77+
78+
try {
79+
this.state = "closed";
80+
this.controller.error(err);
81+
} catch {
82+
// Already errored or closed
83+
}
84+
}
85+
86+
close(): void {
87+
if (this.state === "closed") return;
88+
89+
this.state = "closing";
90+
if (this.controller) {
91+
try {
92+
this.controller.close();
93+
} catch {
94+
// Already closed
95+
}
96+
}
97+
this.state = "closed";
98+
}
99+
100+
get closed(): boolean {
101+
return this.state === "closed";
102+
}
103+
104+
get currentState(): SseState {
105+
return this.state;
106+
}
107+
}

‎packages/agent/src/test/fixtures/api.ts‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
import { execFile } from "node:child_process";
2+
import { existsSync } from "node:fs";
3+
import { mkdir, readFile, rm, writeFile } from "node:fs/promises";
4+
import { tmpdir } from "node:os";
5+
import { join } from "node:path";
6+
import { promisify } from "node:util";
7+
import { vi } from "vitest";
8+
import type { PostHogAPIClient } from "../../posthog-api.js";
9+
import type { TaskRun, TreeSnapshot } from "../../types.js";
10+
11+
const execFileAsync = promisify(execFile);
12+
13+
export interface TestRepo {
14+
path: string;
15+
cleanup: () => Promise<void>;
16+
git: (args: string[]) => Promise<string>;
17+
writeFile: (relativePath: string, content: string) => Promise<void>;
18+
readFile: (relativePath: string) => Promise<string>;
19+
deleteFile: (relativePath: string) => Promise<void>;
20+
exists: (relativePath: string) => boolean;
21+
}
22+
23+
export async function createTestRepo(prefix = "test-repo"): Promise<TestRepo> {
24+
const repoPath = join(
25+
tmpdir(),
26+
`${prefix}-${Date.now()}-${Math.random().toString(36).slice(2)}`,
27+
);
28+
await mkdir(repoPath, { recursive: true });
29+
30+
const git = async (args: string[]): Promise<string> => {
31+
const { stdout } = await execFileAsync("git", args, { cwd: repoPath });
32+
return stdout.trim();
33+
};
34+
35+
await git(["init"]);
36+
await git(["config", "user.email", "test@test.com"]);
37+
await git(["config", "user.name", "Test"]);
38+
39+
await writeFile(join(repoPath, ".gitignore"), ".posthog/\n");
40+
await writeFile(join(repoPath, "README.md"), "# Test Repo");
41+
await git(["add", "."]);
42+
await git(["commit", "-m", "Initial commit"]);
43+
44+
return {
45+
path: repoPath,
46+
cleanup: () => rm(repoPath, { recursive: true, force: true }),
47+
git,
48+
writeFile: async (relativePath: string, content: string) => {
49+
const fullPath = join(repoPath, relativePath);
50+
const dir = join(fullPath, "..");
51+
await mkdir(dir, { recursive: true });
52+
await writeFile(fullPath, content);
53+
},
54+
readFile: async (relativePath: string) => {
55+
return readFile(join(repoPath, relativePath), "utf-8");
56+
},
57+
deleteFile: async (relativePath: string) => {
58+
await rm(join(repoPath, relativePath), { force: true });
59+
},
60+
exists: (relativePath: string) => {
61+
return existsSync(join(repoPath, relativePath));
62+
},
63+
};
64+
}
65+
66+
export function createMockApiClient(
67+
overrides: Partial<PostHogAPIClient> = {},
68+
): PostHogAPIClient {
69+
return {
70+
uploadTaskArtifacts: vi
71+
.fn()
72+
.mockResolvedValue([{ storage_path: "gs://bucket/trees/test.tar.gz" }]),
73+
downloadArtifact: vi.fn(),
74+
getTaskRun: vi.fn(),
75+
fetchTaskRunLogs: vi.fn(),
76+
...overrides,
77+
} as unknown as PostHogAPIClient;
78+
}
79+
80+
export function createTaskRun(overrides: Partial<TaskRun> = {}): TaskRun {
81+
return {
82+
id: "run-1",
83+
task: "task-1",
84+
team: 1,
85+
branch: null,
86+
stage: null,
87+
environment: "local",
88+
status: "in_progress",
89+
log_url: "https://logs.example.com/run-1",
90+
error_message: null,
91+
output: null,
92+
state: {},
93+
created_at: new Date().toISOString(),
94+
updated_at: new Date().toISOString(),
95+
completed_at: null,
96+
...overrides,
97+
};
98+
}
99+
100+
export function createSnapshot(
101+
overrides: Partial<TreeSnapshot> = {},
102+
): TreeSnapshot {
103+
return {
104+
treeHash: "test-tree-hash",
105+
baseCommit: null,
106+
archiveUrl: "gs://bucket/trees/test.tar.gz",
107+
changes: [],
108+
timestamp: new Date().toISOString(),
109+
...overrides,
110+
};
111+
}

‎packages/agent/src/test/fixtures/config.ts‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
import type { AgentServerConfig } from "../../server/types.js";
2+
import type { TestRepo } from "./api.js";
3+
4+
export type { AgentServerConfig };
5+
6+
export function createAgentServerConfig(
7+
repo: TestRepo,
8+
overrides: Partial<AgentServerConfig> = {},
9+
): AgentServerConfig {
10+
return {
11+
port: 3001,
12+
repositoryPath: repo.path,
13+
apiUrl: "http://localhost:8000",
14+
apiKey: "test-api-key",
15+
projectId: 1,
16+
jwtSecret: "test-jwt-secret",
17+
...overrides,
18+
};
19+
}

0 commit comments

Comments
 (0)