z-ai-web-dev-sdk/dist/index.test.js

SDK for Z AI Web Dev

import fs from 'fs/promises';
import ZAI from './index';
// Mock the fs/promises module
jest.mock('fs/promises');
const mockedFs = fs;
// Mock the global fetch function
global.fetch = jest.fn();
const mockedFetch = global.fetch;
describe('ZAI SDK', () => {
    const mockConfig = {
        baseUrl: 'https://api.example.com',
        apiKey: 'test-api-key',
        chatId: 'test-chat-id',
        userId: 'test-user-id',
    };
    const mockChatBody = {
        messages: [{ role: 'user', content: 'Hello' }],
    };
    const mockChatBodyWithModel = {
        ...mockChatBody,
        model: 'test-model',
    };
    const mockApiResponse = {
        id: 'chatcmpl-123',
        object: 'chat.completion',
        created: 1677652288,
        model: 'test-model',
        choices: [
            {
                index: 0,
                message: {
                    role: 'assistant',
                    content: 'Hi there!',
                },
                finish_reason: 'stop',
            },
        ],
    };
    const mockSearchArgs = {
        query: 'TypeScript SDK',
        num: 5,
        recency_days: 10,
    };
    const mockSearchResults = [
        {
            url: 'https://example.com/typescript-sdk',
            name: 'TypeScript SDK Guide',
            snippet: 'Learn how to use TypeScript SDK effectively',
            host_name: 'example.com',
            rank: 1,
            date: '2024-01-15',
            favicon: 'https://example.com/favicon.ico',
        },
        {
            url: 'https://docs.example.com/sdk',
            name: 'Official SDK Documentation',
            snippet: 'Comprehensive guide for SDK integration',
            host_name: 'docs.example.com',
            rank: 2,
            date: '2024-01-14',
            favicon: 'https://docs.example.com/favicon.ico',
        },
    ];
    const mockFunctionResponse = {
        result: mockSearchResults,
    };
    beforeEach(() => {
        // Reset mocks before each test
        mockedFs.readFile.mockClear();
        mockedFetch.mockClear();
    });
    it('should load config and create a ZAI instance', async () => {
        mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConfig));
        const zai = await ZAI.create();
        expect(zai).toBeInstanceOf(ZAI);
        // Verify that it tried to read the config file
        expect(mockedFs.readFile).toHaveBeenCalled();
    });
    it('should throw an error if config file is not found', async () => {
        mockedFs.readFile.mockRejectedValue({ code: 'ENOENT' });
        await expect(ZAI.create()).rejects.toThrow('Configuration file not found or invalid.');
    });
    it('should make a successful chat completion request', async () => {
        mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConfig));
        mockedFetch.mockResolvedValue({
            ok: true,
            json: () => Promise.resolve(mockApiResponse),
        });
        const zai = await ZAI.create();
        const completion = await zai.chat.completions.create(mockChatBodyWithModel);
        expect(completion).toEqual(mockApiResponse);
        expect(mockedFetch).toHaveBeenCalledWith(`${mockConfig.baseUrl}/chat/completions`, {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
                'Authorization': `Bearer ${mockConfig.apiKey}`,
                'X-Z-AI-From': 'Z',
                'X-Chat-Id': mockConfig.chatId,
                'X-User-Id': mockConfig.userId,
            },
            body: JSON.stringify(mockChatBodyWithModel),
        });
    });
    it('should make a successful chat completion request without a model', async () => {
        mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConfig));
        mockedFetch.mockResolvedValue({
            ok: true,
            json: () => Promise.resolve(mockApiResponse),
        });
        const zai = await ZAI.create();
        const completion = await zai.chat.completions.create(mockChatBody);
        expect(completion).toEqual(mockApiResponse);
        expect(mockedFetch).toHaveBeenCalledWith(`${mockConfig.baseUrl}/chat/completions`, {
            method: 'POST',
            headers: {
                'Content-Type': 'application/json',
                'Authorization': `Bearer ${mockConfig.apiKey}`,
                'X-Z-AI-From': 'Z',
                'X-Chat-Id': mockConfig.chatId,
                'X-User-Id': mockConfig.userId,
            },
            body: JSON.stringify(mockChatBody),
        });
    });
    it('should handle API errors during chat completion', async () => {
        const errorResponse = { status: 500, statusText: 'Internal Server Error' };
        mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConfig));
        mockedFetch.mockResolvedValue({
            ok: false,
            status: errorResponse.status,
            text: () => Promise.resolve(errorResponse.statusText),
        });
        const zai = await ZAI.create();
        await expect(zai.chat.completions.create(mockChatBodyWithModel)).rejects.toThrow(`API request failed with status ${errorResponse.status}: ${errorResponse.statusText}`);
    });
    describe('Function execution capabilities', () => {
        beforeEach(() => {
            mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConfig));
        });
        it('should make a successful search function call', async () => {
            mockedFetch.mockResolvedValue({
                ok: true,
                json: () => Promise.resolve(mockFunctionResponse),
            });
            const zai = await ZAI.create();
            const results = await zai.functions.invoke('web_search', mockSearchArgs);
            expect(results).toEqual(mockSearchResults);
            expect(mockedFetch).toHaveBeenCalledWith(`${mockConfig.baseUrl}/functions/invoke`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'Authorization': `Bearer ${mockConfig.apiKey}`,
                    'X-Z-AI-From': 'Z',
                },
                body: JSON.stringify({
                    function_name: 'search',
                    arguments: mockSearchArgs,
                }),
            });
        });
        it('should make a search function call with minimal arguments', async () => {
            const minimalArgs = { query: 'test query' };
            mockedFetch.mockResolvedValue({
                ok: true,
                json: () => Promise.resolve(mockFunctionResponse),
            });
            const zai = await ZAI.create();
            const results = await zai.functions.invoke('web_search', minimalArgs);
            expect(results).toEqual(mockSearchResults);
            expect(mockedFetch).toHaveBeenCalledWith(`${mockConfig.baseUrl}/functions/invoke`, {
                method: 'POST',
                headers: {
                    'Content-Type': 'application/json',
                    'Authorization': `Bearer ${mockConfig.apiKey}`,
                    'X-Z-AI-From': 'Z',
                },
                body: JSON.stringify({
                    function_name: 'search',
                    arguments: minimalArgs,
                }),
            });
        });
        it('should handle function execution errors', async () => {
            const errorResponse = { status: 400, statusText: 'Bad Request' };
            mockedFetch.mockResolvedValue({
                ok: false,
                status: errorResponse.status,
                text: () => Promise.resolve(errorResponse.statusText),
            });
            const zai = await ZAI.create();
            await expect(zai.functions.invoke('web_search', mockSearchArgs)).rejects.toThrow(`Function invoke failed with status ${errorResponse.status}: ${errorResponse.statusText}`);
        });
        it('should handle network errors during function execution', async () => {
            mockedFetch.mockRejectedValue(new Error('Network error'));
            const zai = await ZAI.create();
            await expect(zai.functions.invoke('web_search', mockSearchArgs)).rejects.toThrow('Network error');
        });
    });
});