Spaces:
Paused
Paused
File size: 4,475 Bytes
fae625a 2df2e7e fae625a 2df2e7e fae625a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 | import { ObjectId } from "mongodb";
import { collections, ready } from "$lib/server/database";
import { config } from "$lib/server/config";
import { sha256 } from "$lib/utils/sha256";
import type { User } from "$lib/types/User";
import type { Session } from "$lib/types/Session";
import type { Conversation } from "$lib/types/Conversation";
import type { Workspace } from "$lib/types/Workspace";
export function createTestLocals(overrides?: Partial<App.Locals>): App.Locals {
return {
sessionId: "test-session-id",
isAdmin: false,
user: undefined,
token: undefined,
...overrides,
};
}
export interface TestUser {
user: User;
session: Session;
locals: App.Locals;
/**
* The raw session secret. `authenticateRequest()` reads this from the cookie and looks
* the session up by its sha256, so it is *not* `session.sessionId`.
*/
secretSessionId: string;
/** Ready-made `Cookie` header value, for driving real authentication. */
cookie: string;
}
export async function createTestUser(): Promise<TestUser> {
await ready;
const userId = new ObjectId();
const secretSessionId = crypto.randomUUID();
const sessionId = await sha256(secretSessionId);
const user: User = {
_id: userId,
createdAt: new Date(),
updatedAt: new Date(),
username: `user-${userId.toString().slice(0, 8)}`,
name: "Test User",
avatarUrl: "https://example.com/avatar.png",
hfUserId: `hf-${userId.toString()}`,
};
const session: Session = {
_id: new ObjectId(),
createdAt: new Date(),
updatedAt: new Date(),
userId,
sessionId,
expiresAt: new Date(Date.now() + 1000 * 60 * 60 * 24),
};
await collections.users.insertOne(user);
await collections.sessions.insertOne(session);
return {
user,
session,
locals: {
user,
sessionId,
isAdmin: false,
token: undefined,
},
secretSessionId,
cookie: `${config.COOKIE_NAME}=${secretSessionId}`,
};
}
export async function createTestConversation(
locals: App.Locals,
overrides?: Partial<Conversation>
): Promise<Conversation> {
await ready;
const conv: Conversation = {
_id: new ObjectId(),
title: "Test Conversation",
model: "test-model",
mode: "chat",
messages: [],
createdAt: new Date(),
updatedAt: new Date(),
...(locals.user ? { userId: locals.user._id } : { sessionId: locals.sessionId }),
...overrides,
};
await collections.conversations.insertOne(conv);
return conv;
}
export async function createTestWorkspace(
locals: App.Locals,
conversationId: ObjectId,
overrides?: Partial<Workspace>
): Promise<Workspace> {
await ready;
const flavor = overrides?.flavor ?? "cpu-basic";
const workspace: Workspace = {
_id: new ObjectId(),
conversationId,
repo: { kind: "blank" },
flavor,
model: "test-model",
thinkingLevel: "medium",
policy: "auto",
sandbox: { status: "unprovisioned", flavor },
spend: { accruedUsd: 0 },
createdAt: new Date(),
updatedAt: new Date(),
...(locals.user ? { userId: locals.user._id } : { sessionId: locals.sessionId }),
...overrides,
};
await collections.workspaces.insertOne(workspace);
return workspace;
}
/** Wipes every collection the app writes to. Anything added to `getCollections()` belongs here. */
export async function cleanupTestData() {
await ready;
await Promise.all([
collections.conversations.deleteMany({}),
collections.workspaces.deleteMany({}),
collections.agentConfigurationLedgers.deleteMany({}),
collections.agentConfigurationReceipts.deleteMany({}),
collections.conversationStats.deleteMany({}),
collections.abortedGenerations.deleteMany({}),
collections.generationEvents.deleteMany({}),
collections.generations.deleteMany({}),
collections.users.deleteMany({}),
collections.sessions.deleteMany({}),
collections.settings.deleteMany({}),
collections.sharedConversations.deleteMany({}),
collections.reports.deleteMany({}),
collections.assistants.deleteMany({}),
collections.messageEvents.deleteMany({}),
collections.semaphores.deleteMany({}),
collections.migrationResults.deleteMany({}),
collections.tokenCaches.deleteMany({}),
collections.tools.deleteMany({}),
cleanupGridFS(),
]);
}
/**
* Deletes files rather than calling `bucket.drop()`, which would take the bucket's indexes
* with it and rejects outright when the namespace was never created.
*/
async function cleanupGridFS() {
const files = await collections.bucket.find({}).toArray();
await Promise.all(files.map((file) => collections.bucket.delete(file._id)));
}
|