@lobehub/chat
Version:
Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.
241 lines (200 loc) • 8.56 kB
text/typescript
// @vitest-environment node
import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import * as desensitizeTool from '../utils/desensitizeUrl';
import {
CloudflareStreamTransformer,
desensitizeCloudflareUrl,
fillUrl,
} from './cloudflareHelpers';
afterEach(() => {
vi.restoreAllMocks();
});
describe('cloudflareHelpers', () => {
describe('CloudflareStreamTransformer', () => {
let transformer: CloudflareStreamTransformer;
beforeEach(() => {
transformer = new CloudflareStreamTransformer();
});
describe('parseChunk', () => {
let chunks: string[];
let controller: TransformStreamDefaultController;
beforeEach(() => {
chunks = [];
controller = Object.create(TransformStreamDefaultController.prototype);
vi.spyOn(controller, 'enqueue').mockImplementation((chunk) => {
chunks.push(chunk);
});
});
it('should parse chunk', () => {
// Arrange
const chunk = 'data: {"key": "value", "response": "response1"}';
const textDecoder = new TextDecoder();
// Act
transformer['parseChunk'](chunk, controller);
// Assert
expect(chunks.length).toBe(2);
expect(chunks[0]).toBe('event: text\n');
expect(chunks[1]).toBe('data: "response1"\n\n');
});
it('should not replace `data` in text', () => {
// Arrange
const chunk = 'data: {"key": "value", "response": "data: a"}';
const textDecoder = new TextDecoder();
// Act
transformer['parseChunk'](chunk, controller);
// Assert
expect(chunks.length).toBe(2);
expect(chunks[0]).toBe('event: text\n');
expect(chunks[1]).toBe('data: "data: a"\n\n');
});
});
describe('transform', () => {
const textDecoder = new TextDecoder();
const textEncoder = new TextEncoder();
let chunks: string[];
beforeEach(() => {
chunks = [];
vi.spyOn(
transformer as any as {
parseChunk: (chunk: string, controller: TransformStreamDefaultController) => void;
},
'parseChunk',
).mockImplementation((chunk: string, _) => {
chunks.push(chunk);
});
});
it('should split single chunk', async () => {
// Arrange
const chunk = textEncoder.encode('data: {"key": "value", "response": "response1"}\n\n');
// Act
await transformer.transform(chunk, undefined!);
// Assert
expect(chunks.length).toBe(1);
expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
});
it('should split multiple chunks', async () => {
// Arrange
const chunk = textEncoder.encode(
'data: {"key": "value", "response": "response1"}\n\n' +
'data: {"key": "value", "response": "response2"}\n\n',
);
// Act
await transformer.transform(chunk, undefined!);
// Assert
expect(chunks.length).toBe(2);
expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
expect(chunks[1]).toBe('data: {"key": "value", "response": "response2"}');
});
it('should ignore empty chunk', async () => {
// Arrange
const chunk = textEncoder.encode('\n\n');
// Act
await transformer.transform(chunk, undefined!);
// Assert
expect(chunks.join()).toBe('');
});
it('should split and concat delayed chunks', async () => {
// Arrange
const chunk1 = textEncoder.encode('data: {"key": "value", "respo');
const chunk2 = textEncoder.encode('nse": "response1"}\n\ndata: {"key": "val');
const chunk3 = textEncoder.encode('ue", "response": "response2"}\n\n');
// Act & Assert
await transformer.transform(chunk1, undefined!);
expect(transformer['parseChunk']).not.toHaveBeenCalled();
expect(chunks.length).toBe(0);
expect(transformer['buffer']).toBe('data: {"key": "value", "respo');
await transformer.transform(chunk2, undefined!);
expect(chunks.length).toBe(1);
expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
expect(transformer['buffer']).toBe('data: {"key": "val');
await transformer.transform(chunk3, undefined!);
expect(chunks.length).toBe(2);
expect(chunks[1]).toBe('data: {"key": "value", "response": "response2"}');
expect(transformer['buffer']).toBe('');
});
it('should ignore standalone [DONE]', async () => {
// Arrange
const chunk = textEncoder.encode('data: [DONE]\n\n');
// Act
await transformer.transform(chunk, undefined!);
// Assert
expect(transformer['parseChunk']).not.toHaveBeenCalled();
expect(chunks.length).toBe(0);
expect(transformer['buffer']).toBe('');
});
it('should ignore [DONE] in chunk', async () => {
// Arrange
const chunk = textEncoder.encode(
'data: {"key": "value", "response": "response1"}\n\ndata: [DONE]\n\n',
);
// Act
await transformer.transform(chunk, undefined!);
// Assert
expect(chunks.length).toBe(1);
expect(chunks[0]).toBe('data: {"key": "value", "response": "response1"}');
expect(transformer['buffer']).toBe('');
});
});
});
describe('fillUrl', () => {
it('should return URL with account id', () => {
const url = fillUrl('80009000a000b000c000d000e000f000');
expect(url).toBe(
'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/',
);
});
});
describe('maskAccountId', () => {
describe('desensitizeAccountId', () => {
it('should replace account id with **** in official API endpoint', () => {
const url =
'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe('https://api.cloudflare.com/client/v4/accounts/****/ai/run/');
});
it('should replace account id with **** in custom API endpoint', () => {
const url =
'https://api.cloudflare.com/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe('https://api.cloudflare.com/custom/prefix/****/custom/suffix/');
});
});
describe('desensitizeCloudflareUrl', () => {
it('should mask account id in official API endpoint', () => {
const url =
'https://api.cloudflare.com/client/v4/accounts/80009000a000b000c000d000e000f000/ai/run/';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe('https://api.cloudflare.com/client/v4/accounts/****/ai/run/');
});
it('should call desensitizeUrl for custom API endpoint', () => {
const url = 'https://custom.url/path';
vi.spyOn(desensitizeTool, 'desensitizeUrl').mockImplementation(
(_) => 'https://custom.mocked.url',
);
const maskedUrl = desensitizeCloudflareUrl(url);
expect(desensitizeTool.desensitizeUrl).toHaveBeenCalledWith('https://custom.url');
expect(maskedUrl).toBe('https://custom.mocked.url/path');
});
it('should mask account id in custom API endpoint', () => {
const url =
'https://custom.url/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe('https://cu****om.url/custom/prefix/****/custom/suffix/');
});
it('should mask account id in custom API endpoint with query params', () => {
const url =
'https://custom.url/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/?query=param';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe(
'https://cu****om.url/custom/prefix/****/custom/suffix/?query=param',
);
});
it('should mask account id in custom API endpoint with port', () => {
const url =
'https://custom.url:8080/custom/prefix/80009000a000b000c000d000e000f000/custom/suffix/';
const maskedUrl = desensitizeCloudflareUrl(url);
expect(maskedUrl).toBe('https://cu****om.url:****/custom/prefix/****/custom/suffix/');
});
});
});
});