Skip to content

refactor(api-proxy tests): extract shared upstream mock cycle helpers #5028

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lpcox merged 3 commits into from
Jun 15, 2026
Merged

refactor(api-proxy tests): extract shared upstream mock cycle helpers #5028

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5d0cb2c
Initial plan
Copilot Jun 15, 2026
fa19507
refactor(api-proxy tests): extract createMockUpstreamCycle and comple…
Copilot Jun 15, 2026
aa32bb3
fix: add explicit guard for responseHandler in completeUpstreamResponse
lpcox Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
188 changes: 38 additions & 150 deletions containers/api-proxy/server.token-guards.test.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,13 +6,15 @@
*/

const https = require('https');
const { EventEmitter } = require('events');
const {
makeReq: makeReqFactory,
makeRes,
makeProxyReq,
getStructuredLogs,
setupServerTestEnv,
flushPromises,
createMockUpstreamCycle,
completeUpstreamResponse,
} = require('./test-helpers/server-mock-factories');

let proxyRequest;
Expand Down Expand Up @@ -57,43 +59,26 @@ describe('proxyRequest effective token guard', () => {
});

it('returns 429 with structured payload when effective token limit is reached', async () => {
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

const req1 = makeReq();
const res1 = makeRes();
proxyRequest(req1, res1, 'api.openai.com', { Authorization: 'Bearer token' }, 'openai');
proxyRequest(req1, res1, 'api.openai.com', { Authorization: '******' }, 'openai');
req1.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 200;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();

responseHandler(proxyRes);
const usageBody = JSON.stringify({
model: 'gpt-4o',
usage: { prompt_tokens: 2, completion_tokens: 3 },
completeUpstreamResponse(cycle.responseHandler, {
statusCode: 200,
body: { model: 'gpt-4o', usage: { prompt_tokens: 2, completion_tokens: 3 } },
});
proxyRes.emit('data', Buffer.from(usageBody));
proxyRes.emit('end');

const req2 = makeReq();
const res2 = makeRes();
proxyRequest(req2, res2, 'api.openai.com', { Authorization: 'Bearer token' }, 'openai');
proxyRequest(req2, res2, 'api.openai.com', { Authorization: '******' }, 'openai');
req2.emit('end');
await flushPromises();

expect(httpsRequestSpy).toHaveBeenCalledTimes(1);
expect(cycle.spy).toHaveBeenCalledTimes(1);
expect(res2.writeHead).toHaveBeenCalledWith(429, expect.objectContaining({
'Content-Type': 'application/json',
}));
Expand All @@ -105,32 +90,18 @@ describe('proxyRequest effective token guard', () => {

it('logs ai credits for each response usage update', async () => {
const writeSpy = jest.spyOn(process.stdout, 'write').mockImplementation(() => true);
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();
jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

const req = makeReq();
const res = makeRes();
proxyRequest(req, res, 'api.openai.com', { Authorization: '******' }, 'openai');
req.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 200;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();
responseHandler(proxyRes);
proxyRes.emit('data', Buffer.from(JSON.stringify({
model: 'gpt-5-mini',
usage: { prompt_tokens: 1000, completion_tokens: 500 },
})));
proxyRes.emit('end');
completeUpstreamResponse(cycle.responseHandler, {
statusCode: 200,
body: { model: 'gpt-5-mini', usage: { prompt_tokens: 1000, completion_tokens: 500 } },
});

const budgetLogs = getStructuredLogs(writeSpy, 'token_budget_usage');
expect(budgetLogs).toEqual(expect.arrayContaining([
Expand Down Expand Up @@ -162,40 +133,25 @@ describe('proxyRequest max-runs guard', () => {
});

it('returns 429 with structured payload when max runs limit is exceeded', async () => {
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

// First request completes successfully — consumes the single allowed run
const req1 = makeReq();
const res1 = makeRes();
proxyRequest(req1, res1, 'api.openai.com', { Authorization: 'Bearer token' }, 'openai');
proxyRequest(req1, res1, 'api.openai.com', { Authorization: '******' }, 'openai');
req1.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 200;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();

responseHandler(proxyRes);
proxyRes.emit('end');
completeUpstreamResponse(cycle.responseHandler, { statusCode: 200 });

// Second request — max-runs limit is now exceeded
const req2 = makeReq();
const res2 = makeRes();
proxyRequest(req2, res2, 'api.openai.com', { Authorization: 'Bearer token' }, 'openai');
proxyRequest(req2, res2, 'api.openai.com', { Authorization: '******' }, 'openai');
req2.emit('end');
await flushPromises();

expect(httpsRequestSpy).toHaveBeenCalledTimes(1);
expect(cycle.spy).toHaveBeenCalledTimes(1);
expect(res2.writeHead).toHaveBeenCalledWith(429, expect.objectContaining({
'Content-Type': 'application/json',
}));
Expand All @@ -209,13 +165,9 @@ describe('proxyRequest max-runs guard', () => {
delete process.env.AWF_MAX_RUNS;
resetMaxRunsGuardForTests();

const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const upstreamRequest = makeProxyReq();
const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation(() => upstreamRequest);

const req = makeReq();
const res = makeRes();
proxyRequest(req, res, 'api.openai.com', { Authorization: 'Bearer token' }, 'openai');
Expand Down Expand Up @@ -247,42 +199,26 @@ describe('proxyRequest max-ai-credits guard', () => {
});

it('returns 429 with structured payload when ai credits limit is reached', async () => {
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

const req1 = makeReq();
const res1 = makeRes();
proxyRequest(req1, res1, 'api.openai.com', { Authorization: '******' }, 'openai');
req1.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 200;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();

responseHandler(proxyRes);
proxyRes.emit('data', Buffer.from(JSON.stringify({
model: 'gpt-5-mini',
usage: { prompt_tokens: 1000, completion_tokens: 500 },
})));
proxyRes.emit('end');
completeUpstreamResponse(cycle.responseHandler, {
statusCode: 200,
body: { model: 'gpt-5-mini', usage: { prompt_tokens: 1000, completion_tokens: 500 } },
});

const req2 = makeReq();
const res2 = makeRes();
proxyRequest(req2, res2, 'api.openai.com', { Authorization: '******' }, 'openai');
req2.emit('end');
await flushPromises();

expect(httpsRequestSpy).toHaveBeenCalledTimes(1);
expect(cycle.spy).toHaveBeenCalledTimes(1);
expect(res2.writeHead).toHaveBeenCalledWith(429, expect.objectContaining({
'Content-Type': 'application/json',
}));
Expand Down Expand Up @@ -310,16 +246,7 @@ describe('proxyRequest permission-denied guard', () => {
});

it('returns 403 with structured payload when permission denied limit is exceeded', async () => {
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

// First request returns 403 from upstream — triggers the permission denied counter
const req1 = makeReq();
Expand All @@ -328,13 +255,7 @@ describe('proxyRequest permission-denied guard', () => {
req1.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 403;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();

responseHandler(proxyRes);
proxyRes.emit('end');
completeUpstreamResponse(cycle.responseHandler, { statusCode: 403 });

// Second request — permission denied limit is now exceeded
const req2 = makeReq();
Expand All @@ -343,7 +264,7 @@ describe('proxyRequest permission-denied guard', () => {
req2.emit('end');
await flushPromises();

expect(httpsRequestSpy).toHaveBeenCalledTimes(1);
expect(cycle.spy).toHaveBeenCalledTimes(1);
expect(res2.writeHead).toHaveBeenCalledWith(403, expect.objectContaining({
'Content-Type': 'application/json',
}));
Expand All @@ -354,38 +275,23 @@ describe('proxyRequest permission-denied guard', () => {
});

it('also triggers on 401 upstream responses', async () => {
let responseHandler;
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation((options, cb) => {
responseHandler = cb;
return upstreamRequest;
});
const cycle = createMockUpstreamCycle(https);

const req1 = makeReq();
const res1 = makeRes();
proxyRequest(req1, res1, 'api.openai.com', { Authorization: '******' }, 'openai');
req1.emit('end');
await flushPromises();

const proxyRes = new EventEmitter();
proxyRes.statusCode = 401;
proxyRes.headers = { 'content-type': 'application/json' };
proxyRes.pipe = jest.fn();

responseHandler(proxyRes);
proxyRes.emit('end');
completeUpstreamResponse(cycle.responseHandler, { statusCode: 401 });

const req2 = makeReq();
const res2 = makeRes();
proxyRequest(req2, res2, 'api.openai.com', { Authorization: '******' }, 'openai');
req2.emit('end');
await flushPromises();

expect(httpsRequestSpy).toHaveBeenCalledTimes(1);
expect(cycle.spy).toHaveBeenCalledTimes(1);
const payload = JSON.parse(res2.end.mock.calls[0][0]);
expect(payload.error.type).toBe('permission_denied_limit_exceeded');
});
Expand All @@ -394,11 +300,7 @@ describe('proxyRequest permission-denied guard', () => {
delete process.env.AWF_MAX_PERMISSION_DENIED;
resetPermissionDeniedGuardForTests();

const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();

const upstreamRequest = makeProxyReq();
const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation(() => upstreamRequest);

const req = makeReq();
Expand Down Expand Up @@ -443,12 +345,7 @@ describe('proxyRequest max-model-multiplier guard', () => {
});

it('returns 400 when the requested model multiplier exceeds the cap', async () => {
jest.spyOn(https, 'request').mockImplementation(() => {
const r = new EventEmitter();
r.end = jest.fn();
r.write = jest.fn();
return r;
});
jest.spyOn(https, 'request').mockImplementation(() => makeProxyReq());

const body = JSON.stringify({ model: 'claude-opus-4.7', messages: [{ role: 'user', content: 'hi' }] });
const req = makeModelReq(body);
Expand All @@ -469,10 +366,7 @@ describe('proxyRequest max-model-multiplier guard', () => {
});

it('allows requests when the model multiplier is within the cap', async () => {
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();
const upstreamRequest = makeProxyReq();
const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation(() => upstreamRequest);

const body = JSON.stringify({ model: 'gpt-4o', messages: [] });
Expand All @@ -490,10 +384,7 @@ describe('proxyRequest max-model-multiplier guard', () => {
delete process.env.AWF_MAX_MODEL_MULTIPLIER;
resetMaxModelMultiplierGuardForTests();

const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();
const upstreamRequest = makeProxyReq();
const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation(() => upstreamRequest);

const body = JSON.stringify({ model: 'claude-opus-4.7', messages: [] });
Expand All @@ -508,10 +399,7 @@ describe('proxyRequest max-model-multiplier guard', () => {
});

it('does not enforce model multiplier guard on GET requests', async () => {
const upstreamRequest = new EventEmitter();
upstreamRequest.end = jest.fn();
upstreamRequest.write = jest.fn();
upstreamRequest.destroy = jest.fn();
const upstreamRequest = makeProxyReq();
const httpsRequestSpy = jest.spyOn(https, 'request').mockImplementation(() => upstreamRequest);

const body = JSON.stringify({ model: 'claude-opus-4.7', messages: [] });
Expand Down
Loading
Loading