134 lines
3.4 KiB
TypeScript
134 lines
3.4 KiB
TypeScript
import { beforeEach, describe, expect, mock, test } from "bun:test";
|
|
import { Hono } from "hono";
|
|
import type { AutumnContext, HonoEnv } from "@/honoUtils/HonoEnv.js";
|
|
|
|
const mockState = {
|
|
getTokenTrackParamsCalls: [] as Record<string, unknown>[],
|
|
runTrackWithRolloutCalls: [] as Record<string, unknown>[],
|
|
queuedForReplay: false,
|
|
};
|
|
|
|
const trackBody = {
|
|
customer_id: "cus_123",
|
|
entity_id: "ent_123",
|
|
feature_id: "ai_credits",
|
|
value: 3.5,
|
|
};
|
|
|
|
const featureDeductions = [
|
|
{
|
|
feature: { id: "ai_credits" },
|
|
deduction: 1,
|
|
tokens: {
|
|
usage: {
|
|
modelName: "openai/gpt-4.1",
|
|
inputTokens: 100,
|
|
outputTokens: 50,
|
|
},
|
|
cost: 3.5,
|
|
},
|
|
},
|
|
];
|
|
|
|
mock.module("@/internal/balances/track/utils/getTokenTrackParams.js", () => ({
|
|
getTokenTrackParams: async (args: Record<string, unknown>) => {
|
|
mockState.getTokenTrackParamsCalls.push(args);
|
|
return { body: trackBody, featureDeductions };
|
|
},
|
|
}));
|
|
|
|
mock.module("@/internal/balances/track/runTrackWithRollout.js", () => ({
|
|
runTrackWithRollout: async (args: {
|
|
ctx: AutumnContext;
|
|
body: typeof trackBody;
|
|
featureDeductions: typeof featureDeductions;
|
|
}) => {
|
|
mockState.runTrackWithRolloutCalls.push(args);
|
|
if (mockState.queuedForReplay) {
|
|
args.ctx.extraLogs.trackQueuedForReplay = true;
|
|
}
|
|
return {
|
|
customer_id: args.body.customer_id,
|
|
entity_id: args.body.entity_id,
|
|
value: args.body.value,
|
|
balance: null,
|
|
};
|
|
},
|
|
}));
|
|
|
|
import { handleTrackTokens } from "@/internal/balances/handlers/handleTrackTokens.js";
|
|
|
|
const requestBody = {
|
|
customer_id: "cus_123",
|
|
entity_id: "ent_123",
|
|
model_id: "openai/gpt-4.1",
|
|
input_tokens: 100,
|
|
output_tokens: 50,
|
|
};
|
|
|
|
const createApp = ({ ctx }: { ctx: AutumnContext }) => {
|
|
const app = new Hono<HonoEnv>();
|
|
app.use("*", async (c, next) => {
|
|
c.set("ctx", ctx);
|
|
await next();
|
|
});
|
|
app.post("/track_tokens", ...handleTrackTokens);
|
|
return app;
|
|
};
|
|
|
|
const createCtx = (): AutumnContext =>
|
|
({
|
|
features: [],
|
|
extraLogs: {},
|
|
scopes: [],
|
|
skipCache: false,
|
|
}) as unknown as AutumnContext;
|
|
|
|
describe("handleTrackTokens", () => {
|
|
beforeEach(() => {
|
|
mockState.getTokenTrackParamsCalls = [];
|
|
mockState.runTrackWithRolloutCalls = [];
|
|
mockState.queuedForReplay = false;
|
|
});
|
|
|
|
test("tracks converted token usage through the rollout path", async () => {
|
|
const ctx = createCtx();
|
|
const response = await createApp({ ctx }).request("/track_tokens", {
|
|
method: "POST",
|
|
headers: { "content-type": "application/json" },
|
|
body: JSON.stringify(requestBody),
|
|
});
|
|
|
|
expect(response.status).toBe(200);
|
|
expect(await response.json()).toEqual({
|
|
customer_id: "cus_123",
|
|
entity_id: "ent_123",
|
|
value: 3.5,
|
|
balance: null,
|
|
});
|
|
expect(mockState.getTokenTrackParamsCalls).toHaveLength(1);
|
|
expect(mockState.getTokenTrackParamsCalls[0]).toMatchObject({
|
|
input: requestBody,
|
|
});
|
|
expect(mockState.runTrackWithRolloutCalls).toHaveLength(1);
|
|
expect(mockState.runTrackWithRolloutCalls[0]).toMatchObject({
|
|
body: trackBody,
|
|
featureDeductions,
|
|
});
|
|
});
|
|
|
|
test("returns 202 when rollout fallback queues token tracking for replay", async () => {
|
|
mockState.queuedForReplay = true;
|
|
const ctx = createCtx();
|
|
const response = await createApp({ ctx }).request("/track_tokens", {
|
|
method: "POST",
|
|
headers: { "content-type": "application/json" },
|
|
body: JSON.stringify(requestBody),
|
|
});
|
|
|
|
expect(response.status).toBe(202);
|
|
expect(ctx.extraLogs.trackQueuedForReplay).toBe(true);
|
|
expect(mockState.runTrackWithRolloutCalls).toHaveLength(1);
|
|
});
|
|
});
|