121 lines
3.0 KiB
TypeScript
121 lines
3.0 KiB
TypeScript
import { assertEquals } from 'jsr:@std/assert@^1.0.0';
|
|
import { ChatHandler } from '../ts/api/handlers/chat.ts';
|
|
import { EmbeddingsHandler } from '../ts/api/handlers/embeddings.ts';
|
|
import { UpstreamTimeoutError } from '../ts/containers/base-container.ts';
|
|
|
|
class TestResponse {
|
|
public statusCode = 200;
|
|
public headers: Record<string, string> = {};
|
|
public body = '';
|
|
|
|
public writeHead(statusCode: number, headers: Record<string, string>): TestResponse {
|
|
this.statusCode = statusCode;
|
|
this.headers = headers;
|
|
return this;
|
|
}
|
|
|
|
public end(body = ''): TestResponse {
|
|
this.body = body;
|
|
return this;
|
|
}
|
|
|
|
public write(_chunk: string | Uint8Array): boolean {
|
|
return true;
|
|
}
|
|
}
|
|
|
|
Deno.test('ChatHandler maps upstream timeouts to 504 responses', async () => {
|
|
const handler = new ChatHandler(
|
|
{
|
|
async findContainerForModel() {
|
|
return {
|
|
async chatCompletion() {
|
|
throw new UpstreamTimeoutError();
|
|
},
|
|
async chatCompletionStream() {
|
|
throw new UpstreamTimeoutError();
|
|
},
|
|
};
|
|
},
|
|
} as never,
|
|
{
|
|
async getModel(modelName: string) {
|
|
return { id: modelName };
|
|
},
|
|
} as never,
|
|
{
|
|
async loadModel() {
|
|
return { success: false };
|
|
},
|
|
} as never,
|
|
{
|
|
shouldDeployLocallyFirst() {
|
|
return false;
|
|
},
|
|
} as never,
|
|
);
|
|
|
|
const response = new TestResponse();
|
|
await handler.handleChatCompletion(
|
|
{ headers: {} } as never,
|
|
response as never,
|
|
{ model: 'meta-llama/Llama-3.1-8B-Instruct', messages: [{ role: 'user', content: 'hi' }] },
|
|
);
|
|
|
|
assertEquals(response.statusCode, 504);
|
|
assertEquals(JSON.parse(response.body).error.type, 'upstream_timeout');
|
|
});
|
|
|
|
Deno.test('EmbeddingsHandler maps upstream timeouts to 504 responses', async () => {
|
|
const originalFetch = globalThis.fetch;
|
|
globalThis.fetch = async () => {
|
|
const error = new Error('request aborted');
|
|
error.name = 'AbortError';
|
|
throw error;
|
|
};
|
|
|
|
try {
|
|
const handler = new EmbeddingsHandler(
|
|
{
|
|
async findContainerForModel() {
|
|
return null;
|
|
},
|
|
} as never,
|
|
{
|
|
async getModel(modelName: string) {
|
|
return { id: modelName };
|
|
},
|
|
} as never,
|
|
{
|
|
async ensureModelViaControlPlane(modelName: string) {
|
|
return {
|
|
location: {
|
|
modelId: modelName,
|
|
nodeName: 'worker-a',
|
|
endpoint: 'http://worker-a:8080',
|
|
healthy: true,
|
|
engine: 'vllm',
|
|
containerId: 'remote',
|
|
},
|
|
};
|
|
},
|
|
getLocalNodeName() {
|
|
return 'control';
|
|
},
|
|
} as never,
|
|
);
|
|
|
|
const response = new TestResponse();
|
|
await handler.handleEmbeddings(
|
|
{ headers: {} } as never,
|
|
response as never,
|
|
{ model: 'BAAI/bge-m3', input: 'hello' },
|
|
);
|
|
|
|
assertEquals(response.statusCode, 504);
|
|
assertEquals(JSON.parse(response.body).error.type, 'upstream_timeout');
|
|
} finally {
|
|
globalThis.fetch = originalFetch;
|
|
}
|
|
});
|