1
import { afterEach, describe, test, expect, jest } from '@jest/globals';
2
import { once } from 'node:events';
3
import { PassThrough } from 'node:stream';
4
import { Response } from 'node-fetch';
5
import { CHAT_COMPLETION_SOURCES } from '../src/constants';
6
import { flattenSchema, forwardFetchResponse } from '../src/util';
7

8
function createMockExpressResponse() {
9
    const response = new PassThrough();
10
    response.statusCode = 200;
11
    response.statusMessage = '';
12

13
    return response;
14
}
15

16
async function collectResponseBody(response) {
17
    const chunks = [];
18

19
    response.on('data', chunk => chunks.push(Buffer.from(chunk)));
20

21
    await once(response, 'finish');
22

23
    return Buffer.concat(chunks).toString('utf8');
24
}
25

26
afterEach(() => {
27
    jest.restoreAllMocks();
28
});
29

30
describe('flattenSchema', () => {
31
    test('should return the schema if it is not an object', () => {
32
        const schema = 'it is not an object';
33
        expect(flattenSchema(schema, CHAT_COMPLETION_SOURCES.MAKERSUITE)).toBe(schema);
34
    });
35

36
    test('should handle schema with $defs and $ref', () => {
37
        const schema = {
38
            $schema: 'http://json-schema.org/draft-07/schema#',
39
            $defs: {
40
                a: { type: 'string' },
41
                b: {
42
                    type: 'object',
43
                    properties: {
44
                        c: { $ref: '#/$defs/a' },
45
                    },
46
                },
47
            },
48
            properties: {
49
                d: { $ref: '#/$defs/b' },
50
            },
51
        };
52
        const expected = {
53
            properties: {
54
                d: {
55
                    type: 'object',
56
                    properties: {
57
                        c: { type: 'string' },
58
                    },
59
                },
60
            },
61
        };
62
        expect(flattenSchema(schema, CHAT_COMPLETION_SOURCES.MAKERSUITE)).toEqual(expected);
63
    });
64

65
    test('should filter unsupported properties for Google API schema', () => {
66
        const schema = {
67
            $defs: {
68
                a: {
69
                    type: 'string',
70
                    default: 'test',
71
                },
72
            },
73
            type: 'object',
74
            properties: {
75
                b: { $ref: '#/$defs/a' },
76
                c: { type: 'number' },
77
            },
78
            additionalProperties: false,
79
            exclusiveMinimum: 0,
80
            propertyNames: {
81
                pattern: '^[A-Za-z_][A-Za-z0-9_]*$',
82
            },
83
        };
84
        const expected = {
85
            type: 'object',
86
            properties: {
87
                b: {
88
                    type: 'string',
89
                },
90
                c: { type: 'number' },
91
            },
92
        };
93
        expect(flattenSchema(schema, CHAT_COMPLETION_SOURCES.MAKERSUITE)).toEqual(expected);
94
    });
95

96
    test('should not filter properties for non-Google API schema', () => {
97
        const schema = {
98
            $defs: {
99
                a: {
100
                    type: 'string',
101
                    default: 'test',
102
                },
103
            },
104
            type: 'object',
105
            properties: {
106
                b: { $ref: '#/$defs/a' },
107
                c: { type: 'number' },
108
            },
109
            additionalProperties: false,
110
            exclusiveMinimum: 0,
111
            propertyNames: {
112
                pattern: '^[A-Za-z_][A-Za-z0-9_]*$',
113
            },
114
        };
115
        const expected = {
116
            type: 'object',
117
            properties: {
118
                b: {
119
                    type: 'string',
120
                    default: 'test',
121
                },
122
                c: { type: 'number' },
123
            },
124
            additionalProperties: false,
125
            exclusiveMinimum: 0,
126
            propertyNames: {
127
                pattern: '^[A-Za-z_][A-Za-z0-9_]*$',
128
            },
129
        };
130
        expect(flattenSchema(schema, 'some-other-api')).toEqual(expected);
131
    });
132
});
133

134
describe('forwardFetchResponse', () => {
135
    test('should log JSON error bodies and return the original body for non-2xx streaming responses', async () => {
136
        const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => undefined);
137
        const body = JSON.stringify({ error: { message: 'Forbidden by upstream policy' }, detail: 'policy_denied' });
138
        const response = createMockExpressResponse();
139
        const bodyPromise = collectResponseBody(response);
140

141
        await forwardFetchResponse(new Response(body, {
142
            status: 403,
143
            statusText: 'Forbidden',
144
        }), response);
145

146
        expect(await bodyPromise).toBe(body);
147
        expect(response.statusCode).toBe(403);
148
        expect(warnSpy).toHaveBeenCalledWith(`Streaming request failed with status 403 Forbidden: ${body}`);
149
    });
150

151
    test('should log plain text error bodies and return the original body for non-2xx streaming responses', async () => {
152
        const warnSpy = jest.spyOn(console, 'warn').mockImplementation(() => undefined);
153
        const body = 'Plain text upstream failure';
154
        const response = createMockExpressResponse();
155
        const bodyPromise = collectResponseBody(response);
156

157
        await forwardFetchResponse(new Response(body, {
158
            status: 502,
159
            statusText: 'Bad Gateway',
160
        }), response);
161

162
        expect(await bodyPromise).toBe(body);
163
        expect(response.statusCode).toBe(502);
164
        expect(warnSpy).toHaveBeenCalledWith(`Streaming request failed with status 502 Bad Gateway: ${body}`);
165
    });
166
});