diff --git a/containers/api-proxy/server.token-guards.test.js b/containers/api-proxy/server.token-guards.test.js index 0c23273dd..df93fcbae 100644 --- a/containers/api-proxy/server.token-guards.test.js +++ b/containers/api-proxy/server.token-guards.test.js @@ -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; @@ -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', })); @@ -105,15 +90,7 @@ 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(); @@ -121,16 +98,10 @@ describe('proxyRequest effective token guard', () => { 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([ @@ -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', })); @@ -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'); @@ -247,16 +199,7 @@ 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(); @@ -264,17 +207,10 @@ describe('proxyRequest max-ai-credits guard', () => { 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(); @@ -282,7 +218,7 @@ describe('proxyRequest max-ai-credits guard', () => { 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', })); @@ -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(); @@ -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(); @@ -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', })); @@ -354,16 +275,7 @@ 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(); @@ -371,13 +283,7 @@ describe('proxyRequest permission-denied guard', () => { 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(); @@ -385,7 +291,7 @@ describe('proxyRequest permission-denied guard', () => { 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'); }); @@ -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(); @@ -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); @@ -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: [] }); @@ -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: [] }); @@ -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: [] }); diff --git a/containers/api-proxy/server.token-steering.test.js b/containers/api-proxy/server.token-steering.test.js index 46c1c4521..20144f2c4 100644 --- a/containers/api-proxy/server.token-steering.test.js +++ b/containers/api-proxy/server.token-steering.test.js @@ -7,7 +7,7 @@ const https = require('https'); const { EventEmitter } = require('events'); -const { setupServerTestEnv, flushPromises } = require('./test-helpers/server-mock-factories'); +const { setupServerTestEnv, flushPromises, makeProxyReq, completeUpstreamResponse } = require('./test-helpers/server-mock-factories'); let proxyRequest; let getAndClearPendingSteeringMessage; @@ -92,15 +92,8 @@ describe('token steering — getAndClearPendingSteeringMessage and injectSteerin const nowSpy = jest.spyOn(Date, 'now').mockReturnValue(start); resetTimeoutSteeringForTests(); - const upstreamReq1 = new EventEmitter(); - upstreamReq1.end = jest.fn(); - upstreamReq1.write = jest.fn(); - upstreamReq1.destroy = jest.fn(); - - const upstreamReq2 = new EventEmitter(); - upstreamReq2.end = jest.fn(); - upstreamReq2.write = jest.fn(); - upstreamReq2.destroy = jest.fn(); + const upstreamReq1 = makeProxyReq(); + const upstreamReq2 = makeProxyReq(); jest.spyOn(https, 'request') .mockImplementationOnce(() => upstreamReq1) @@ -144,20 +137,9 @@ describe('token steering — getAndClearPendingSteeringMessage and injectSteerin it('injects 80% warning into an OpenAI request body and clears it on the next request', async () => { // Two upstream request objects — one per proxyRequest call. let responseHandler; - const upstreamReq1 = new EventEmitter(); - upstreamReq1.end = jest.fn(); - upstreamReq1.write = jest.fn(); - upstreamReq1.destroy = jest.fn(); - - const upstreamReq2 = new EventEmitter(); - upstreamReq2.end = jest.fn(); - upstreamReq2.write = jest.fn(); - upstreamReq2.destroy = jest.fn(); - - const upstreamReq3 = new EventEmitter(); - upstreamReq3.end = jest.fn(); - upstreamReq3.write = jest.fn(); - upstreamReq3.destroy = jest.fn(); + const upstreamReq1 = makeProxyReq(); + const upstreamReq2 = makeProxyReq(); + const upstreamReq3 = makeProxyReq(); jest.spyOn(https, 'request') .mockImplementationOnce((_opts, cb) => { responseHandler = cb; return upstreamReq1; }) @@ -175,17 +157,10 @@ describe('token steering — getAndClearPendingSteeringMessage and injectSteerin 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-4o', - usage: { prompt_tokens: 0, completion_tokens: 21 }, - }))); - proxyRes.emit('end'); + completeUpstreamResponse(responseHandler, { + statusCode: 200, + body: { model: 'gpt-4o', usage: { prompt_tokens: 0, completion_tokens: 21 } }, + }); // Request 2: the proxy should inject the 80% warning into the outgoing body. // We send a minimal OpenAI chat body and inspect what the proxy writes upstream. @@ -245,15 +220,8 @@ describe('token steering — getAndClearPendingSteeringMessage and injectSteerin const { proxyRequest: localProxyRequest } = require('./server'); let responseHandler; - const upstreamReq1 = new EventEmitter(); - upstreamReq1.end = jest.fn(); - upstreamReq1.write = jest.fn(); - upstreamReq1.destroy = jest.fn(); - - const upstreamReq2 = new EventEmitter(); - upstreamReq2.end = jest.fn(); - upstreamReq2.write = jest.fn(); - upstreamReq2.destroy = jest.fn(); + const upstreamReq1 = makeProxyReq(); + const upstreamReq2 = makeProxyReq(); jest.spyOn(https, 'request') .mockImplementationOnce((_opts, cb) => { responseHandler = cb; return upstreamReq1; }) @@ -270,17 +238,10 @@ describe('token steering — getAndClearPendingSteeringMessage and injectSteerin 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-4o', - usage: { prompt_tokens: 0, completion_tokens: 21 }, - }))); - proxyRes.emit('end'); + completeUpstreamResponse(responseHandler, { + statusCode: 200, + body: { model: 'gpt-4o', usage: { prompt_tokens: 0, completion_tokens: 21 } }, + }); // Request 2: steering is disabled, so no warning should be injected. const req2Body = Buffer.from(JSON.stringify({ diff --git a/containers/api-proxy/test-helpers/server-mock-factories.js b/containers/api-proxy/test-helpers/server-mock-factories.js index 97b37b577..1efcd436d 100644 --- a/containers/api-proxy/test-helpers/server-mock-factories.js +++ b/containers/api-proxy/test-helpers/server-mock-factories.js @@ -84,6 +84,58 @@ function flushPromises() { return new Promise(resolve => setImmediate(resolve)); } +/** + * Creates a mock upstream HTTP request cycle for testing proxyRequest guards. + * + * Sets up a `jest.spyOn` on `https.request` that captures the upstream response + * handler callback and returns a pre-built mock upstream request. After + * `flushPromises()` has drained the async body-collection chain, callers can + * access the captured callback via the `responseHandler` getter. + * + * @param {object} https - The `https` module to spy on (imported by the test file). + * @returns {{ spy: jest.SpyInstance, upstreamRequest: EventEmitter, responseHandler: Function|undefined }} + */ +function createMockUpstreamCycle(https) { + let responseHandler; + const upstreamRequest = makeProxyReq(); + + const spy = jest.spyOn(https, 'request').mockImplementation((_options, cb) => { + responseHandler = cb; + return upstreamRequest; + }); + + return { + spy, + upstreamRequest, + /** The upstream response handler captured when https.request was called. */ + get responseHandler() { return responseHandler; }, + }; +} + +/** + * Completes an upstream response cycle: invokes the captured response handler + * with a mocked proxy response, optionally emits a data chunk containing a + * JSON-serialised body, then emits `'end'`. + * + * @param {Function} responseHandler - The upstream response callback (e.g. `cycle.responseHandler`). + * @param {object} [options] + * @param {number} [options.statusCode=200] - HTTP status code of the upstream response. + * @param {object|null} [options.body=null] - JSON-serialisable payload to emit as a data chunk. + * @returns {EventEmitter} The mocked proxy response emitter. + */ +function completeUpstreamResponse(responseHandler, { statusCode = 200, body = null } = {}) { + if (typeof responseHandler !== 'function') { + throw new Error('completeUpstreamResponse: responseHandler is not a function — ensure the https.request callback has been captured before calling this helper'); + } + const proxyRes = makeProxyRes(statusCode); + responseHandler(proxyRes); + if (body !== null) { + proxyRes.emit('data', Buffer.from(JSON.stringify(body))); + } + proxyRes.emit('end'); + return proxyRes; +} + module.exports = { makeReq, makeRes, @@ -92,4 +144,6 @@ module.exports = { getStructuredLogs, setupServerTestEnv, flushPromises, + createMockUpstreamCycle, + completeUpstreamResponse, };