UNPKG

@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
// @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/'); }); }); }); });