pricing-no-price.test.ts 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246
  1. import { beforeEach, describe, expect, it, vi } from "vitest";
  2. import type { SystemSettings } from "@/types/system-config";
  3. const { cloudSyncRequests, requestCloudPriceTableSyncMock } = vi.hoisted(() => {
  4. const cloudSyncRequests: Array<{ reason: string }> = [];
  5. const requestCloudPriceTableSyncMock = vi.fn((payload: { reason: string }) => {
  6. cloudSyncRequests.push(payload);
  7. });
  8. return { cloudSyncRequests, requestCloudPriceTableSyncMock };
  9. });
  10. vi.mock("@/lib/price-sync/cloud-price-updater", () => ({
  11. requestCloudPriceTableSync: requestCloudPriceTableSyncMock,
  12. }));
  13. vi.mock("@/lib/logger", () => ({
  14. logger: {
  15. debug: vi.fn(),
  16. info: vi.fn(),
  17. warn: vi.fn(),
  18. error: vi.fn(),
  19. trace: vi.fn(),
  20. },
  21. }));
  22. vi.mock("@/repository/model-price", () => ({
  23. findLatestPriceByModel: vi.fn(),
  24. }));
  25. vi.mock("@/repository/system-config", () => ({
  26. getSystemSettings: vi.fn(),
  27. }));
  28. vi.mock("@/repository/message", () => ({
  29. updateMessageRequestCost: vi.fn(),
  30. updateMessageRequestDetails: vi.fn(),
  31. updateMessageRequestDuration: vi.fn(),
  32. }));
  33. vi.mock("@/lib/session-manager", () => ({
  34. SessionManager: {
  35. updateSessionUsage: vi.fn(async () => {}),
  36. storeSessionResponse: vi.fn(async () => {}),
  37. extractCodexPromptCacheKey: vi.fn(),
  38. updateSessionWithCodexCacheKey: vi.fn(async () => {}),
  39. },
  40. }));
  41. vi.mock("@/lib/rate-limit", () => ({
  42. RateLimitService: {
  43. trackCost: vi.fn(),
  44. trackUserDailyCost: vi.fn(),
  45. },
  46. }));
  47. vi.mock("@/lib/proxy-status-tracker", () => ({
  48. ProxyStatusTracker: {
  49. getInstance: () => ({
  50. endRequest: vi.fn(),
  51. }),
  52. },
  53. }));
  54. import { finalizeRequestStats } from "@/app/v1/_lib/proxy/response-handler";
  55. import { ProxySession } from "@/app/v1/_lib/proxy/session";
  56. import { RateLimitService } from "@/lib/rate-limit";
  57. import { updateMessageRequestCost } from "@/repository/message";
  58. import { findLatestPriceByModel } from "@/repository/model-price";
  59. import { getSystemSettings } from "@/repository/system-config";
  60. function makeSystemSettings(
  61. billingModelSource: SystemSettings["billingModelSource"]
  62. ): SystemSettings {
  63. const now = new Date();
  64. return {
  65. id: 1,
  66. siteTitle: "test",
  67. allowGlobalUsageView: false,
  68. currencyDisplay: "USD",
  69. billingModelSource,
  70. timezone: null,
  71. enableAutoCleanup: false,
  72. cleanupRetentionDays: 30,
  73. cleanupSchedule: "0 2 * * *",
  74. cleanupBatchSize: 10000,
  75. enableClientVersionCheck: false,
  76. verboseProviderError: false,
  77. enableHttp2: false,
  78. interceptAnthropicWarmupRequests: false,
  79. enableResponseFixer: true,
  80. responseFixerConfig: {
  81. fixTruncatedJson: true,
  82. fixSseFormat: true,
  83. fixEncoding: true,
  84. maxJsonDepth: 200,
  85. maxFixSize: 1024 * 1024,
  86. },
  87. createdAt: now,
  88. updatedAt: now,
  89. };
  90. }
  91. function createSession({
  92. originalModel,
  93. redirectedModel,
  94. }: {
  95. originalModel: string;
  96. redirectedModel: string;
  97. }): ProxySession {
  98. const session = new (
  99. ProxySession as unknown as {
  100. new (init: {
  101. startTime: number;
  102. method: string;
  103. requestUrl: URL;
  104. headers: Headers;
  105. headerLog: string;
  106. request: { message: Record<string, unknown>; log: string; model: string | null };
  107. userAgent: string | null;
  108. context: unknown;
  109. clientAbortSignal: AbortSignal | null;
  110. }): ProxySession;
  111. }
  112. )({
  113. startTime: Date.now(),
  114. method: "POST",
  115. requestUrl: new URL("http://localhost/v1/messages"),
  116. headers: new Headers(),
  117. headerLog: "",
  118. request: { message: {}, log: "(test)", model: redirectedModel },
  119. userAgent: null,
  120. context: {},
  121. clientAbortSignal: null,
  122. });
  123. session.setOriginalModel(originalModel);
  124. session.setSessionId("sess-test");
  125. const provider = {
  126. id: 99,
  127. name: "test-provider",
  128. providerType: "claude",
  129. costMultiplier: 1.0,
  130. streamingIdleTimeoutMs: 0,
  131. } as any;
  132. const user = {
  133. id: 123,
  134. name: "test-user",
  135. dailyResetTime: "00:00",
  136. dailyResetMode: "fixed",
  137. } as any;
  138. const key = {
  139. id: 456,
  140. name: "test-key",
  141. dailyResetTime: "00:00",
  142. dailyResetMode: "fixed",
  143. } as any;
  144. session.setProvider(provider);
  145. session.setAuthState({
  146. user,
  147. key,
  148. apiKey: "sk-test",
  149. success: true,
  150. });
  151. session.setMessageContext({
  152. id: 2000,
  153. createdAt: new Date(),
  154. user,
  155. key,
  156. apiKey: "sk-test",
  157. });
  158. return session;
  159. }
  160. describe("价格表缺失/查询失败:请求不计费且不报错", () => {
  161. beforeEach(() => {
  162. cloudSyncRequests.splice(0, cloudSyncRequests.length);
  163. vi.clearAllMocks();
  164. });
  165. it("无价格:应跳过 DB cost 更新与限流 cost 追踪,并触发异步同步", async () => {
  166. vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));
  167. vi.mocked(findLatestPriceByModel).mockResolvedValue(null);
  168. const session = createSession({ originalModel: "m1", redirectedModel: "m2" });
  169. const responseText = JSON.stringify({
  170. type: "message",
  171. usage: { input_tokens: 2, output_tokens: 3 },
  172. });
  173. await finalizeRequestStats(session, responseText, 200, 5);
  174. expect(updateMessageRequestCost).not.toHaveBeenCalled();
  175. expect(RateLimitService.trackCost).not.toHaveBeenCalled();
  176. expect(findLatestPriceByModel).toHaveBeenCalled();
  177. expect(cloudSyncRequests).toEqual([{ reason: "missing-model" }]);
  178. });
  179. it("价格数据为空对象:应视为无价格并触发异步同步", async () => {
  180. vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("redirected"));
  181. vi.mocked(findLatestPriceByModel).mockImplementation(async (modelName: string) => {
  182. if (modelName === "m2") {
  183. return {
  184. id: 1,
  185. modelName: "m2",
  186. priceData: {},
  187. source: "litellm",
  188. createdAt: new Date(),
  189. updatedAt: new Date(),
  190. } as any;
  191. }
  192. return null;
  193. });
  194. const session = createSession({ originalModel: "m1", redirectedModel: "m2" });
  195. const responseText = JSON.stringify({
  196. type: "message",
  197. usage: { input_tokens: 2, output_tokens: 3 },
  198. });
  199. await finalizeRequestStats(session, responseText, 200, 5);
  200. expect(updateMessageRequestCost).not.toHaveBeenCalled();
  201. expect(RateLimitService.trackCost).not.toHaveBeenCalled();
  202. expect(cloudSyncRequests).toEqual([{ reason: "missing-model" }]);
  203. });
  204. it("价格查询抛错:应跳过计费且不影响响应", async () => {
  205. vi.mocked(getSystemSettings).mockResolvedValue(makeSystemSettings("original"));
  206. vi.mocked(findLatestPriceByModel).mockImplementation(async () => {
  207. throw new Error("db query failed");
  208. });
  209. const session = createSession({ originalModel: "m1", redirectedModel: "m2" });
  210. const responseText = JSON.stringify({
  211. type: "message",
  212. usage: { input_tokens: 2, output_tokens: 3 },
  213. });
  214. await finalizeRequestStats(session, responseText, 200, 5);
  215. expect(updateMessageRequestCost).not.toHaveBeenCalled();
  216. expect(RateLimitService.trackCost).not.toHaveBeenCalled();
  217. });
  218. });