This is page 6 of 7. Use http://codebase.md/stripe/agent-toolkit?lines=true&page={x} to view the full context.
# Directory Structure
```
├── .github
│ ├── ISSUE_TEMPLATE
│ │ ├── bug_report.yml
│ │ ├── config.yml
│ │ └── feature_request.yml
│ └── workflows
│ ├── main.yml
│ ├── npm_release_shared.yml
│ ├── pypi_release.yml
│ └── sync-skills.yml
├── .gitignore
├── .vscode
│ ├── extensions.json
│ ├── launch.json
│ └── settings.json
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── gemini-extension.json
├── LICENSE
├── llm
│ ├── ai-sdk
│ │ ├── jest.config.ts
│ │ ├── LICENSE
│ │ ├── meter
│ │ │ ├── examples
│ │ │ │ ├── .env.example
│ │ │ │ ├── .gitignore
│ │ │ │ ├── anthropic.ts
│ │ │ │ ├── google.ts
│ │ │ │ ├── openai.ts
│ │ │ │ ├── README.md
│ │ │ │ └── tsconfig.json
│ │ │ ├── index.ts
│ │ │ ├── meter-event-logging.ts
│ │ │ ├── meter-event-types.ts
│ │ │ ├── README.md
│ │ │ ├── tests
│ │ │ │ ├── ai-sdk-billing-wrapper-anthropic.test.ts
│ │ │ │ ├── ai-sdk-billing-wrapper-general.test.ts
│ │ │ │ ├── ai-sdk-billing-wrapper-google.test.ts
│ │ │ │ ├── ai-sdk-billing-wrapper-openai.test.ts
│ │ │ │ ├── ai-sdk-billing-wrapper-other-providers.test.ts
│ │ │ │ ├── meter-event-logging.test.ts
│ │ │ │ └── model-name-normalization.test.ts
│ │ │ ├── tsconfig.json
│ │ │ ├── types.ts
│ │ │ ├── utils.ts
│ │ │ └── wrapperV2.ts
│ │ ├── package.json
│ │ ├── pnpm-lock.yaml
│ │ ├── provider
│ │ │ ├── examples
│ │ │ │ ├── .env.example
│ │ │ │ ├── .gitignore
│ │ │ │ ├── anthropic.ts
│ │ │ │ ├── google.ts
│ │ │ │ ├── openai.ts
│ │ │ │ ├── README.md
│ │ │ │ └── tsconfig.json
│ │ │ ├── index.ts
│ │ │ ├── README.md
│ │ │ ├── stripe-language-model.ts
│ │ │ ├── stripe-provider.ts
│ │ │ ├── tests
│ │ │ │ ├── stripe-language-model.test.ts
│ │ │ │ ├── stripe-provider.test.ts
│ │ │ │ └── utils.test.ts
│ │ │ ├── tsconfig.build.json
│ │ │ ├── tsconfig.json
│ │ │ ├── types.ts
│ │ │ └── utils.ts
│ │ ├── README.md
│ │ ├── tsconfig.json
│ │ └── tsup.config.ts
│ ├── README.md
│ └── token-meter
│ ├── examples
│ │ ├── anthropic.ts
│ │ ├── gemini.ts
│ │ └── openai.ts
│ ├── index.ts
│ ├── jest.config.ts
│ ├── LICENSE
│ ├── meter-event-logging.ts
│ ├── meter-event-types.ts
│ ├── package.json
│ ├── pnpm-lock.yaml
│ ├── README.md
│ ├── tests
│ │ ├── meter-event-logging.test.ts
│ │ ├── model-name-normalization.test.ts
│ │ ├── token-meter-anthropic.test.ts
│ │ ├── token-meter-gemini.test.ts
│ │ ├── token-meter-general.test.ts
│ │ ├── token-meter-openai.test.ts
│ │ └── type-detection.test.ts
│ ├── token-meter.ts
│ ├── tsconfig.build.json
│ ├── tsconfig.json
│ ├── types.ts
│ └── utils
│ └── type-detection.ts
├── README.md
├── SECURITY.md
├── skills
│ ├── get-started-kiro.md
│ ├── README.md
│ ├── stripe-best-practices.md
│ └── sync.js
└── tools
├── modelcontextprotocol
│ ├── .dxtignore
│ ├── .gitignore
│ ├── .node-version
│ ├── .prettierrc
│ ├── build-dxt.js
│ ├── Dockerfile
│ ├── eslint.config.mjs
│ ├── jest.config.ts
│ ├── LICENSE
│ ├── manifest.json
│ ├── package.json
│ ├── pnpm-lock.yaml
│ ├── README.md
│ ├── server.json
│ ├── src
│ │ ├── index.ts
│ │ └── test
│ │ └── index.test.ts
│ ├── stripe_icon.png
│ └── tsconfig.json
├── python
│ ├── .editorconfig
│ ├── .flake8
│ ├── examples
│ │ ├── crewai
│ │ │ ├── .env.template
│ │ │ ├── main.py
│ │ │ └── README.md
│ │ ├── langchain
│ │ │ ├── __init__.py
│ │ │ ├── .env.template
│ │ │ ├── main.py
│ │ │ └── README.md
│ │ ├── openai
│ │ │ ├── .env.template
│ │ │ ├── customer_support
│ │ │ │ ├── .env.template
│ │ │ │ ├── emailer.py
│ │ │ │ ├── env.py
│ │ │ │ ├── main.py
│ │ │ │ ├── pyproject.toml
│ │ │ │ ├── README.md
│ │ │ │ ├── repl.py
│ │ │ │ └── support_agent.py
│ │ │ ├── file_search
│ │ │ │ ├── main.py
│ │ │ │ └── README.md
│ │ │ └── web_search
│ │ │ ├── .env.template
│ │ │ ├── main.py
│ │ │ └── README.md
│ │ └── strands
│ │ └── main.py
│ ├── Makefile
│ ├── pyproject.toml
│ ├── README.md
│ ├── requirements.txt
│ ├── stripe_agent_toolkit
│ │ ├── __init__.py
│ │ ├── api.py
│ │ ├── configuration.py
│ │ ├── crewai
│ │ │ ├── tool.py
│ │ │ └── toolkit.py
│ │ ├── functions.py
│ │ ├── langchain
│ │ │ ├── tool.py
│ │ │ └── toolkit.py
│ │ ├── openai
│ │ │ ├── hooks.py
│ │ │ ├── tool.py
│ │ │ └── toolkit.py
│ │ ├── prompts.py
│ │ ├── schema.py
│ │ ├── strands
│ │ │ ├── __init__.py
│ │ │ ├── hooks.py
│ │ │ ├── tool.py
│ │ │ └── toolkit.py
│ │ └── tools.py
│ └── tests
│ ├── __init__.py
│ ├── test_configuration.py
│ └── test_functions.py
├── README.md
└── typescript
├── .gitignore
├── .prettierrc
├── eslint.config.mjs
├── examples
│ ├── ai-sdk
│ │ ├── .env.template
│ │ ├── index.ts
│ │ ├── package.json
│ │ ├── README.md
│ │ └── tsconfig.json
│ ├── cloudflare
│ │ ├── .dev.vars.example
│ │ ├── .gitignore
│ │ ├── biome.json
│ │ ├── package.json
│ │ ├── README.md
│ │ ├── src
│ │ │ ├── app.ts
│ │ │ ├── imageGenerator.ts
│ │ │ ├── index.ts
│ │ │ ├── oauth.ts
│ │ │ └── utils.ts
│ │ ├── tsconfig.json
│ │ ├── worker-configuration.d.ts
│ │ └── wrangler.jsonc
│ ├── langchain
│ │ ├── .env.template
│ │ ├── index.ts
│ │ ├── package.json
│ │ ├── README.md
│ │ └── tsconfig.json
│ └── openai
│ ├── .env.template
│ ├── index.ts
│ ├── package.json
│ ├── README.md
│ └── tsconfig.json
├── jest.config.ts
├── LICENSE
├── package.json
├── pnpm-lock.yaml
├── pnpm-workspace.yaml
├── README.md
├── src
│ ├── ai-sdk
│ │ ├── index.ts
│ │ ├── tool.ts
│ │ └── toolkit.ts
│ ├── cloudflare
│ │ ├── index.ts
│ │ └── README.md
│ ├── langchain
│ │ ├── index.ts
│ │ ├── tool.ts
│ │ └── toolkit.ts
│ ├── modelcontextprotocol
│ │ ├── index.ts
│ │ ├── README.md
│ │ ├── register-paid-tool.ts
│ │ └── toolkit.ts
│ ├── openai
│ │ ├── index.ts
│ │ └── toolkit.ts
│ ├── shared
│ │ ├── api.ts
│ │ ├── balance
│ │ │ └── retrieveBalance.ts
│ │ ├── configuration.ts
│ │ ├── coupons
│ │ │ ├── createCoupon.ts
│ │ │ └── listCoupons.ts
│ │ ├── customers
│ │ │ ├── createCustomer.ts
│ │ │ └── listCustomers.ts
│ │ ├── disputes
│ │ │ ├── listDisputes.ts
│ │ │ └── updateDispute.ts
│ │ ├── documentation
│ │ │ └── searchDocumentation.ts
│ │ ├── invoiceItems
│ │ │ └── createInvoiceItem.ts
│ │ ├── invoices
│ │ │ ├── createInvoice.ts
│ │ │ ├── finalizeInvoice.ts
│ │ │ └── listInvoices.ts
│ │ ├── paymentIntents
│ │ │ └── listPaymentIntents.ts
│ │ ├── paymentLinks
│ │ │ └── createPaymentLink.ts
│ │ ├── prices
│ │ │ ├── createPrice.ts
│ │ │ └── listPrices.ts
│ │ ├── products
│ │ │ ├── createProduct.ts
│ │ │ └── listProducts.ts
│ │ ├── refunds
│ │ │ └── createRefund.ts
│ │ ├── subscriptions
│ │ │ ├── cancelSubscription.ts
│ │ │ ├── listSubscriptions.ts
│ │ │ └── updateSubscription.ts
│ │ └── tools.ts
│ └── test
│ ├── modelcontextprotocol
│ │ └── register-paid-tool.test.ts
│ └── shared
│ ├── balance
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ ├── configuration.test.ts
│ ├── customers
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ ├── disputes
│ │ └── functions.test.ts
│ ├── documentation
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ ├── invoiceItems
│ │ ├── functions.test.ts
│ │ ├── parameters.test.ts
│ │ └── prompts.test.ts
│ ├── invoices
│ │ ├── functions.test.ts
│ │ ├── parameters.test.ts
│ │ └── prompts.test.ts
│ ├── paymentIntents
│ │ ├── functions.test.ts
│ │ ├── parameters.test.ts
│ │ └── prompts.test.ts
│ ├── paymentLinks
│ │ ├── functions.test.ts
│ │ ├── parameters.test.ts
│ │ └── prompts.test.ts
│ ├── prices
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ ├── products
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ ├── refunds
│ │ ├── functions.test.ts
│ │ └── parameters.test.ts
│ └── subscriptions
│ ├── functions.test.ts
│ ├── parameters.test.ts
│ └── prompts.test.ts
├── tsconfig.json
└── tsup.config.ts
```
# Files
--------------------------------------------------------------------------------
/llm/token-meter/tests/token-meter-openai.test.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Tests for TokenMeter - OpenAI Provider
3 | */
4 |
5 | import Stripe from 'stripe';
6 | import {createTokenMeter} from '../token-meter';
7 | import type {MeterConfig} from '../types';
8 |
9 | // Mock Stripe
10 | jest.mock('stripe');
11 |
12 | describe('TokenMeter - OpenAI Provider', () => {
13 | let mockMeterEventsCreate: jest.Mock;
14 | let config: MeterConfig;
15 | const TEST_API_KEY = 'sk_test_mock_key';
16 |
17 | beforeEach(() => {
18 | jest.clearAllMocks();
19 | mockMeterEventsCreate = jest.fn().mockResolvedValue({});
20 |
21 | // Mock the Stripe constructor
22 | (Stripe as unknown as jest.Mock).mockImplementation(() => ({
23 | v2: {
24 | billing: {
25 | meterEvents: {
26 | create: mockMeterEventsCreate,
27 | },
28 | },
29 | },
30 | }));
31 |
32 | config = {};
33 | });
34 |
35 | describe('Chat Completions - Non-streaming', () => {
36 | it('should track usage from basic chat completion', async () => {
37 | const meter = createTokenMeter(TEST_API_KEY, config);
38 |
39 | const response = {
40 | id: 'chatcmpl-123',
41 | object: 'chat.completion',
42 | created: Date.now(),
43 | model: 'gpt-4o-mini',
44 | choices: [
45 | {
46 | index: 0,
47 | message: {
48 | role: 'assistant',
49 | content: 'Hello, World!',
50 | },
51 | finish_reason: 'stop',
52 | },
53 | ],
54 | usage: {
55 | prompt_tokens: 12,
56 | completion_tokens: 5,
57 | total_tokens: 17,
58 | },
59 | };
60 |
61 | meter.trackUsage(response as any, 'cus_123');
62 |
63 | // Wait for fire-and-forget logging to complete
64 | await new Promise(resolve => setImmediate(resolve));
65 |
66 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
67 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
68 | expect.objectContaining({
69 | event_name: 'token-billing-tokens',
70 | payload: expect.objectContaining({
71 | stripe_customer_id: 'cus_123',
72 | value: '12',
73 | model: 'openai/gpt-4o-mini',
74 | token_type: 'input',
75 | }),
76 | })
77 | );
78 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
79 | expect.objectContaining({
80 | payload: expect.objectContaining({
81 | value: '5',
82 | token_type: 'output',
83 | }),
84 | })
85 | );
86 | });
87 |
88 | it('should track usage from chat completion with tools', async () => {
89 | const meter = createTokenMeter(TEST_API_KEY, config);
90 |
91 | const response = {
92 | id: 'chatcmpl-123',
93 | object: 'chat.completion',
94 | created: Date.now(),
95 | model: 'gpt-4o',
96 | choices: [
97 | {
98 | index: 0,
99 | message: {
100 | role: 'assistant',
101 | content: null,
102 | tool_calls: [
103 | {
104 | id: 'call_123',
105 | type: 'function',
106 | function: {
107 | name: 'get_weather',
108 | arguments: '{"location":"San Francisco"}',
109 | },
110 | },
111 | ],
112 | },
113 | finish_reason: 'tool_calls',
114 | },
115 | ],
116 | usage: {
117 | prompt_tokens: 100,
118 | completion_tokens: 30,
119 | total_tokens: 130,
120 | },
121 | };
122 |
123 | meter.trackUsage(response as any, 'cus_456');
124 |
125 | // Wait for fire-and-forget logging to complete
126 | await new Promise(resolve => setImmediate(resolve));
127 |
128 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
129 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
130 | expect.objectContaining({
131 | payload: expect.objectContaining({
132 | stripe_customer_id: 'cus_456',
133 | value: '100',
134 | model: 'openai/gpt-4o',
135 | token_type: 'input',
136 | }),
137 | })
138 | );
139 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
140 | expect.objectContaining({
141 | payload: expect.objectContaining({
142 | value: '30',
143 | token_type: 'output',
144 | }),
145 | })
146 | );
147 | });
148 |
149 | it('should handle missing usage data', async () => {
150 | const meter = createTokenMeter(TEST_API_KEY, config);
151 |
152 | const response = {
153 | id: 'chatcmpl-123',
154 | object: 'chat.completion',
155 | created: Date.now(),
156 | model: 'gpt-4',
157 | choices: [
158 | {
159 | index: 0,
160 | message: {
161 | role: 'assistant',
162 | content: 'Hello!',
163 | },
164 | finish_reason: 'stop',
165 | },
166 | ],
167 | };
168 |
169 | meter.trackUsage(response as any, 'cus_123');
170 |
171 | // Wait for fire-and-forget logging to complete
172 | await new Promise(resolve => setImmediate(resolve));
173 |
174 | // Should not create events with 0 tokens (code only sends when > 0)
175 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(0);
176 | });
177 |
178 | it('should handle multi-turn conversations', async () => {
179 | const meter = createTokenMeter(TEST_API_KEY, config);
180 |
181 | const response = {
182 | id: 'chatcmpl-789',
183 | object: 'chat.completion',
184 | created: Date.now(),
185 | model: 'gpt-4',
186 | choices: [
187 | {
188 | index: 0,
189 | message: {
190 | role: 'assistant',
191 | content: 'The weather is sunny.',
192 | },
193 | finish_reason: 'stop',
194 | },
195 | ],
196 | usage: {
197 | prompt_tokens: 150, // Includes conversation history
198 | completion_tokens: 10,
199 | total_tokens: 160,
200 | },
201 | };
202 |
203 | meter.trackUsage(response as any, 'cus_123');
204 |
205 | // Wait for fire-and-forget logging to complete
206 | await new Promise(resolve => setImmediate(resolve));
207 |
208 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
209 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
210 | expect.objectContaining({
211 | payload: expect.objectContaining({
212 | value: '150',
213 | model: 'openai/gpt-4',
214 | token_type: 'input',
215 | }),
216 | })
217 | );
218 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
219 | expect.objectContaining({
220 | payload: expect.objectContaining({
221 | value: '10',
222 | token_type: 'output',
223 | }),
224 | })
225 | );
226 | });
227 | });
228 |
229 | describe('Chat Completions - Streaming', () => {
230 | it('should track usage from basic streaming chat', async () => {
231 | const meter = createTokenMeter(TEST_API_KEY, config);
232 |
233 | const chunks = [
234 | {
235 | id: 'chatcmpl-123',
236 | object: 'chat.completion.chunk',
237 | created: Date.now(),
238 | model: 'gpt-4o-mini',
239 | choices: [
240 | {
241 | index: 0,
242 | delta: {content: 'Hello'},
243 | finish_reason: null,
244 | },
245 | ],
246 | },
247 | {
248 | id: 'chatcmpl-123',
249 | object: 'chat.completion.chunk',
250 | created: Date.now(),
251 | model: 'gpt-4o-mini',
252 | choices: [
253 | {
254 | index: 0,
255 | delta: {content: ', World!'},
256 | finish_reason: 'stop',
257 | },
258 | ],
259 | usage: {
260 | prompt_tokens: 12,
261 | completion_tokens: 5,
262 | total_tokens: 17,
263 | },
264 | },
265 | ];
266 |
267 | const mockStream = createMockStreamWithTee(chunks);
268 | const wrappedStream = meter.trackUsageStreamOpenAI(mockStream as any, 'cus_123');
269 |
270 | for await (const _chunk of wrappedStream) {
271 | // Consume stream
272 | }
273 |
274 | // Wait for fire-and-forget logging to complete
275 | await new Promise(resolve => setImmediate(resolve));
276 |
277 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
278 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
279 | expect.objectContaining({
280 | payload: expect.objectContaining({
281 | stripe_customer_id: 'cus_123',
282 | value: '12',
283 | model: 'openai/gpt-4o-mini',
284 | token_type: 'input',
285 | }),
286 | })
287 | );
288 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
289 | expect.objectContaining({
290 | payload: expect.objectContaining({
291 | value: '5',
292 | token_type: 'output',
293 | }),
294 | })
295 | );
296 | });
297 |
298 | it('should track usage from streaming chat with tools', async () => {
299 | const meter = createTokenMeter(TEST_API_KEY, config);
300 |
301 | const chunks = [
302 | {
303 | id: 'chatcmpl-123',
304 | object: 'chat.completion.chunk',
305 | created: Date.now(),
306 | model: 'gpt-4o',
307 | choices: [
308 | {
309 | index: 0,
310 | delta: {
311 | tool_calls: [
312 | {
313 | index: 0,
314 | id: 'call_123',
315 | type: 'function',
316 | function: {
317 | name: 'get_weather',
318 | arguments: '{"location":',
319 | },
320 | },
321 | ],
322 | },
323 | finish_reason: null,
324 | },
325 | ],
326 | },
327 | {
328 | id: 'chatcmpl-123',
329 | object: 'chat.completion.chunk',
330 | created: Date.now(),
331 | model: 'gpt-4o',
332 | choices: [
333 | {
334 | index: 0,
335 | delta: {},
336 | finish_reason: 'tool_calls',
337 | },
338 | ],
339 | usage: {
340 | prompt_tokens: 100,
341 | completion_tokens: 30,
342 | total_tokens: 130,
343 | },
344 | },
345 | ];
346 |
347 | const mockStream = createMockStreamWithTee(chunks);
348 | const wrappedStream = meter.trackUsageStreamOpenAI(mockStream as any, 'cus_456');
349 |
350 | for await (const _chunk of wrappedStream) {
351 | // Consume stream
352 | }
353 |
354 | // Wait for fire-and-forget logging to complete
355 | await new Promise(resolve => setImmediate(resolve));
356 |
357 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
358 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
359 | expect.objectContaining({
360 | payload: expect.objectContaining({
361 | stripe_customer_id: 'cus_456',
362 | value: '100',
363 | model: 'openai/gpt-4o',
364 | token_type: 'input',
365 | }),
366 | })
367 | );
368 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
369 | expect.objectContaining({
370 | payload: expect.objectContaining({
371 | value: '30',
372 | token_type: 'output',
373 | }),
374 | })
375 | );
376 | });
377 |
378 | it('should properly tee the stream', async () => {
379 | const meter = createTokenMeter(TEST_API_KEY, config);
380 |
381 | const chunks = [
382 | {
383 | id: 'chatcmpl-123',
384 | object: 'chat.completion.chunk',
385 | created: Date.now(),
386 | model: 'gpt-4',
387 | choices: [
388 | {
389 | index: 0,
390 | delta: {content: 'Hello'},
391 | finish_reason: null,
392 | },
393 | ],
394 | },
395 | {
396 | id: 'chatcmpl-123',
397 | object: 'chat.completion.chunk',
398 | created: Date.now(),
399 | model: 'gpt-4',
400 | choices: [
401 | {
402 | index: 0,
403 | delta: {content: ' world'},
404 | finish_reason: 'stop',
405 | },
406 | ],
407 | usage: {
408 | prompt_tokens: 10,
409 | completion_tokens: 5,
410 | total_tokens: 15,
411 | },
412 | },
413 | ];
414 |
415 | const mockStream = createMockStreamWithTee(chunks);
416 | const wrappedStream = meter.trackUsageStreamOpenAI(mockStream as any, 'cus_123');
417 |
418 | const receivedChunks: any[] = [];
419 | for await (const chunk of wrappedStream) {
420 | receivedChunks.push(chunk);
421 | }
422 |
423 | expect(receivedChunks).toHaveLength(2);
424 | expect(receivedChunks[0].choices[0].delta.content).toBe('Hello');
425 | expect(receivedChunks[1].choices[0].delta.content).toBe(' world');
426 | });
427 | });
428 |
429 | describe('Responses API - Non-streaming', () => {
430 | it('should track usage from basic responses API', async () => {
431 | const meter = createTokenMeter(TEST_API_KEY, config);
432 |
433 | const response = {
434 | id: 'resp_123',
435 | object: 'response',
436 | created: Date.now(),
437 | model: 'gpt-4o-mini',
438 | output: 'Hello, World!',
439 | usage: {
440 | input_tokens: 15,
441 | output_tokens: 8,
442 | },
443 | };
444 |
445 | meter.trackUsage(response as any, 'cus_123');
446 |
447 | // Wait for fire-and-forget logging to complete
448 | await new Promise(resolve => setImmediate(resolve));
449 |
450 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
451 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
452 | expect.objectContaining({
453 | payload: expect.objectContaining({
454 | stripe_customer_id: 'cus_123',
455 | value: '15',
456 | model: 'openai/gpt-4o-mini',
457 | token_type: 'input',
458 | }),
459 | })
460 | );
461 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
462 | expect.objectContaining({
463 | payload: expect.objectContaining({
464 | value: '8',
465 | token_type: 'output',
466 | }),
467 | })
468 | );
469 | });
470 |
471 | it('should track usage from responses API parse', async () => {
472 | const meter = createTokenMeter(TEST_API_KEY, config);
473 |
474 | const response = {
475 | id: 'resp_456',
476 | object: 'response',
477 | created: Date.now(),
478 | model: 'gpt-4o',
479 | output: {parsed: {city: 'San Francisco', temperature: 72}},
480 | usage: {
481 | input_tokens: 50,
482 | output_tokens: 20,
483 | },
484 | };
485 |
486 | meter.trackUsage(response as any, 'cus_789');
487 |
488 | // Wait for fire-and-forget logging to complete
489 | await new Promise(resolve => setImmediate(resolve));
490 |
491 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
492 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
493 | expect.objectContaining({
494 | payload: expect.objectContaining({
495 | stripe_customer_id: 'cus_789',
496 | value: '50',
497 | model: 'openai/gpt-4o',
498 | token_type: 'input',
499 | }),
500 | })
501 | );
502 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
503 | expect.objectContaining({
504 | payload: expect.objectContaining({
505 | value: '20',
506 | token_type: 'output',
507 | }),
508 | })
509 | );
510 | });
511 | });
512 |
513 | describe('Responses API - Streaming', () => {
514 | it('should track usage from streaming responses API', async () => {
515 | const meter = createTokenMeter(TEST_API_KEY, config);
516 |
517 | const chunks = [
518 | {
519 | type: 'response.output_text.delta',
520 | delta: 'Hello',
521 | },
522 | {
523 | type: 'response.output_text.delta',
524 | delta: ', World!',
525 | },
526 | {
527 | type: 'response.done',
528 | response: {
529 | id: 'resp_123',
530 | model: 'gpt-4o-mini',
531 | usage: {
532 | input_tokens: 15,
533 | output_tokens: 8,
534 | },
535 | },
536 | },
537 | ];
538 |
539 | const mockStream = createMockStreamWithTee(chunks);
540 | const wrappedStream = meter.trackUsageStreamOpenAI(mockStream as any, 'cus_123');
541 |
542 | for await (const _chunk of wrappedStream) {
543 | // Consume stream
544 | }
545 |
546 | // Wait for fire-and-forget logging to complete
547 | await new Promise(resolve => setImmediate(resolve));
548 |
549 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
550 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
551 | expect.objectContaining({
552 | payload: expect.objectContaining({
553 | stripe_customer_id: 'cus_123',
554 | value: '15',
555 | model: 'openai/gpt-4o-mini',
556 | token_type: 'input',
557 | }),
558 | })
559 | );
560 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
561 | expect.objectContaining({
562 | payload: expect.objectContaining({
563 | value: '8',
564 | token_type: 'output',
565 | }),
566 | })
567 | );
568 | });
569 | });
570 |
571 | describe('Embeddings', () => {
572 | it('should track usage from single text embedding', async () => {
573 | const meter = createTokenMeter(TEST_API_KEY, config);
574 |
575 | const response = {
576 | object: 'list',
577 | data: [
578 | {
579 | object: 'embedding',
580 | embedding: new Array(1536).fill(0.1),
581 | index: 0,
582 | },
583 | ],
584 | model: 'text-embedding-ada-002',
585 | usage: {
586 | prompt_tokens: 8,
587 | total_tokens: 8,
588 | },
589 | };
590 |
591 | meter.trackUsage(response as any, 'cus_123');
592 |
593 | // Wait for fire-and-forget logging to complete
594 | await new Promise(resolve => setImmediate(resolve));
595 |
596 | // Embeddings only have input tokens, no output tokens
597 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(1);
598 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
599 | expect.objectContaining({
600 | payload: expect.objectContaining({
601 | stripe_customer_id: 'cus_123',
602 | value: '8',
603 | model: 'openai/text-embedding-ada-002',
604 | token_type: 'input',
605 | }),
606 | })
607 | );
608 | });
609 |
610 | it('should track usage from batch embeddings', async () => {
611 | const meter = createTokenMeter(TEST_API_KEY, config);
612 |
613 | const response = {
614 | object: 'list',
615 | data: [
616 | {
617 | object: 'embedding',
618 | embedding: new Array(1536).fill(0.1),
619 | index: 0,
620 | },
621 | {
622 | object: 'embedding',
623 | embedding: new Array(1536).fill(0.2),
624 | index: 1,
625 | },
626 | {
627 | object: 'embedding',
628 | embedding: new Array(1536).fill(0.3),
629 | index: 2,
630 | },
631 | ],
632 | model: 'text-embedding-3-small',
633 | usage: {
634 | prompt_tokens: 24,
635 | total_tokens: 24,
636 | },
637 | };
638 |
639 | meter.trackUsage(response as any, 'cus_456');
640 |
641 | // Wait for fire-and-forget logging to complete
642 | await new Promise(resolve => setImmediate(resolve));
643 |
644 | // Embeddings only have input tokens, no output tokens
645 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(1);
646 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
647 | expect.objectContaining({
648 | payload: expect.objectContaining({
649 | stripe_customer_id: 'cus_456',
650 | value: '24',
651 | model: 'openai/text-embedding-3-small',
652 | token_type: 'input',
653 | }),
654 | })
655 | );
656 | });
657 |
658 | it('should handle missing usage data in embeddings', async () => {
659 | const meter = createTokenMeter(TEST_API_KEY, config);
660 |
661 | const response = {
662 | object: 'list',
663 | data: [
664 | {
665 | object: 'embedding',
666 | embedding: new Array(1536).fill(0.1),
667 | index: 0,
668 | },
669 | ],
670 | model: 'text-embedding-ada-002',
671 | };
672 |
673 | meter.trackUsage(response as any, 'cus_123');
674 |
675 | // Wait for fire-and-forget logging to complete
676 | await new Promise(resolve => setImmediate(resolve));
677 |
678 | // Should not create events with 0 tokens
679 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(0);
680 | });
681 | });
682 | });
683 |
684 | // Helper function to create mock streams with tee()
685 | function createMockStreamWithTee(chunks: any[]) {
686 | return {
687 | tee() {
688 | const stream1 = {
689 | async *[Symbol.asyncIterator]() {
690 | for (const chunk of chunks) {
691 | yield chunk;
692 | }
693 | },
694 | tee() {
695 | const s1 = {
696 | async *[Symbol.asyncIterator]() {
697 | for (const chunk of chunks) {
698 | yield chunk;
699 | }
700 | },
701 | };
702 | const s2 = {
703 | async *[Symbol.asyncIterator]() {
704 | for (const chunk of chunks) {
705 | yield chunk;
706 | }
707 | },
708 | };
709 | return [s1, s2];
710 | },
711 | };
712 | const stream2 = {
713 | async *[Symbol.asyncIterator]() {
714 | for (const chunk of chunks) {
715 | yield chunk;
716 | }
717 | },
718 | };
719 | return [stream1, stream2];
720 | },
721 | async *[Symbol.asyncIterator]() {
722 | for (const chunk of chunks) {
723 | yield chunk;
724 | }
725 | },
726 | };
727 | }
728 |
729 |
```
--------------------------------------------------------------------------------
/llm/token-meter/tests/token-meter-gemini.test.ts:
--------------------------------------------------------------------------------
```typescript
1 | /**
2 | * Tests for TokenMeter - Gemini Provider
3 | */
4 |
5 | import Stripe from 'stripe';
6 | import {createTokenMeter} from '../token-meter';
7 | import type {MeterConfig} from '../types';
8 |
9 | // Mock Stripe
10 | jest.mock('stripe');
11 |
12 | describe('TokenMeter - Gemini Provider', () => {
13 | let mockMeterEventsCreate: jest.Mock;
14 | let config: MeterConfig;
15 | const TEST_API_KEY = 'sk_test_mock_key';
16 |
17 | beforeEach(() => {
18 | jest.clearAllMocks();
19 | mockMeterEventsCreate = jest.fn().mockResolvedValue({});
20 |
21 | // Mock the Stripe constructor
22 | (Stripe as unknown as jest.Mock).mockImplementation(() => ({
23 | v2: {
24 | billing: {
25 | meterEvents: {
26 | create: mockMeterEventsCreate,
27 | },
28 | },
29 | },
30 | }));
31 |
32 | config = {};
33 | });
34 |
35 | describe('GenerateContent - Non-streaming', () => {
36 | it('should track usage from basic text generation', async () => {
37 | const meter = createTokenMeter(TEST_API_KEY, config);
38 |
39 | const response = {
40 | response: {
41 | text: () => 'Hello, World!',
42 | usageMetadata: {
43 | promptTokenCount: 12,
44 | candidatesTokenCount: 8,
45 | totalTokenCount: 20,
46 | },
47 | modelVersion: 'gemini-2.0-flash-exp',
48 | },
49 | };
50 |
51 | meter.trackUsage(response, 'cus_123');
52 |
53 | // Wait for fire-and-forget logging to complete
54 | await new Promise(resolve => setImmediate(resolve));
55 |
56 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
57 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
58 | expect.objectContaining({
59 | payload: expect.objectContaining({
60 | stripe_customer_id: 'cus_123',
61 | value: '12',
62 | model: 'google/gemini-2.0-flash-exp',
63 | token_type: 'input',
64 | }),
65 | })
66 | );
67 | });
68 |
69 | it('should track usage with reasoning tokens for extended thinking models', async () => {
70 | const meter = createTokenMeter(TEST_API_KEY, config);
71 |
72 | const response = {
73 | response: {
74 | text: () => 'Detailed response after thinking',
75 | usageMetadata: {
76 | promptTokenCount: 20,
77 | candidatesTokenCount: 15,
78 | thoughtsTokenCount: 50, // Reasoning/thinking tokens
79 | totalTokenCount: 85,
80 | },
81 | modelVersion: 'gemini-2.0-flash-thinking-exp',
82 | },
83 | };
84 |
85 | meter.trackUsage(response, 'cus_456');
86 |
87 | // Wait for fire-and-forget logging to complete
88 | await new Promise(resolve => setImmediate(resolve));
89 |
90 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
91 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
92 | expect.objectContaining({
93 | payload: expect.objectContaining({
94 | stripe_customer_id: 'cus_456',
95 | value: '20',
96 | model: 'google/gemini-2.0-flash-thinking-exp',
97 | token_type: 'input',
98 | }),
99 | })
100 | );
101 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
102 | expect.objectContaining({
103 | payload: expect.objectContaining({
104 | value: '65', // 15 candidates + 50 thoughts
105 | token_type: 'output',
106 | }),
107 | })
108 | );
109 | });
110 |
111 | it('should track usage from generation with function calling', async () => {
112 | const meter = createTokenMeter(TEST_API_KEY, config);
113 |
114 | const response = {
115 | response: {
116 | text: () => '',
117 | functionCalls: () => [
118 | {
119 | name: 'get_weather',
120 | args: {location: 'San Francisco'},
121 | },
122 | ],
123 | usageMetadata: {
124 | promptTokenCount: 100,
125 | candidatesTokenCount: 30,
126 | totalTokenCount: 130,
127 | },
128 | modelVersion: 'gemini-1.5-pro',
129 | },
130 | };
131 |
132 | meter.trackUsage(response, 'cus_789');
133 |
134 | // Wait for fire-and-forget logging to complete
135 | await new Promise(resolve => setImmediate(resolve));
136 |
137 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
138 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
139 | expect.objectContaining({
140 | payload: expect.objectContaining({
141 | value: '100',
142 | model: 'google/gemini-1.5-pro',
143 | token_type: 'input',
144 | }),
145 | })
146 | );
147 | });
148 |
149 | it('should track usage with system instructions', async () => {
150 | const meter = createTokenMeter(TEST_API_KEY, config);
151 |
152 | const response = {
153 | response: {
154 | text: () => 'I am following the system instructions.',
155 | usageMetadata: {
156 | promptTokenCount: 50, // Includes system instruction tokens
157 | candidatesTokenCount: 12,
158 | totalTokenCount: 62,
159 | },
160 | modelVersion: 'gemini-2.5-flash',
161 | },
162 | };
163 |
164 | meter.trackUsage(response, 'cus_123');
165 |
166 | // Wait for fire-and-forget logging to complete
167 | await new Promise(resolve => setImmediate(resolve));
168 |
169 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
170 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
171 | expect.objectContaining({
172 | payload: expect.objectContaining({
173 | value: '50',
174 | model: 'google/gemini-2.5-flash',
175 | token_type: 'input',
176 | }),
177 | })
178 | );
179 | });
180 |
181 | it('should use default model name when modelVersion is missing', async () => {
182 | const meter = createTokenMeter(TEST_API_KEY, config);
183 |
184 | const response = {
185 | response: {
186 | text: () => 'Hello',
187 | usageMetadata: {
188 | promptTokenCount: 5,
189 | candidatesTokenCount: 3,
190 | totalTokenCount: 8,
191 | },
192 | },
193 | };
194 |
195 | meter.trackUsage(response, 'cus_999');
196 |
197 | // Wait for fire-and-forget logging to complete
198 | await new Promise(resolve => setImmediate(resolve));
199 |
200 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
201 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
202 | expect.objectContaining({
203 | payload: expect.objectContaining({
204 | value: '5',
205 | model: 'google/gemini',
206 | token_type: 'input',
207 | }),
208 | })
209 | );
210 | });
211 | });
212 |
213 | describe('GenerateContent - Streaming', () => {
214 | it('should require model name parameter', () => {
215 | const meter = createTokenMeter(TEST_API_KEY, config);
216 |
217 | const mockGeminiStream = {
218 | stream: {
219 | async *[Symbol.asyncIterator]() {
220 | yield {
221 | text: () => 'Hello',
222 | usageMetadata: {
223 | promptTokenCount: 10,
224 | candidatesTokenCount: 5,
225 | totalTokenCount: 15,
226 | },
227 | };
228 | },
229 | },
230 | response: Promise.resolve({
231 | text: () => 'Hello',
232 | modelVersion: 'gemini-1.5-pro',
233 | }),
234 | };
235 |
236 | // TypeScript will enforce model name parameter at compile time
237 | // @ts-expect-error - Testing that TypeScript requires model name
238 | meter.trackUsageStreamGemini(mockGeminiStream, 'cus_123');
239 | });
240 |
241 | it('should track usage from basic streaming generation', async () => {
242 | const meter = createTokenMeter(TEST_API_KEY, config);
243 |
244 | const chunks = [
245 | {
246 | text: () => 'Hello',
247 | usageMetadata: null,
248 | },
249 | {
250 | text: () => ', World!',
251 | usageMetadata: {
252 | promptTokenCount: 12,
253 | candidatesTokenCount: 8,
254 | totalTokenCount: 20,
255 | },
256 | },
257 | ];
258 |
259 | const mockGeminiStream = {
260 | stream: {
261 | async *[Symbol.asyncIterator]() {
262 | for (const chunk of chunks) {
263 | yield chunk;
264 | }
265 | },
266 | },
267 | response: Promise.resolve({
268 | text: () => 'Hello, World!',
269 | modelVersion: 'gemini-2.0-flash-exp',
270 | }),
271 | };
272 |
273 | const wrappedStream = meter.trackUsageStreamGemini(
274 | mockGeminiStream,
275 | 'cus_123',
276 | 'gemini-2.0-flash-exp'
277 | );
278 |
279 | for await (const _chunk of wrappedStream.stream) {
280 | // Consume stream
281 | }
282 |
283 | // Wait for fire-and-forget logging to complete
284 | await new Promise(resolve => setImmediate(resolve));
285 |
286 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
287 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
288 | expect.objectContaining({
289 | payload: expect.objectContaining({
290 | stripe_customer_id: 'cus_123',
291 | value: '12',
292 | model: 'google/gemini-2.0-flash-exp',
293 | token_type: 'input',
294 | }),
295 | })
296 | );
297 | });
298 |
299 | it('should track usage from streaming with reasoning tokens', async () => {
300 | const meter = createTokenMeter(TEST_API_KEY, config);
301 |
302 | const chunks = [
303 | {
304 | text: () => 'Thinking...',
305 | usageMetadata: null,
306 | },
307 | {
308 | text: () => 'After consideration, here is my answer.',
309 | usageMetadata: {
310 | promptTokenCount: 20,
311 | candidatesTokenCount: 15,
312 | thoughtsTokenCount: 50,
313 | totalTokenCount: 85,
314 | },
315 | },
316 | ];
317 |
318 | const mockGeminiStream = {
319 | stream: {
320 | async *[Symbol.asyncIterator]() {
321 | for (const chunk of chunks) {
322 | yield chunk;
323 | }
324 | },
325 | },
326 | response: Promise.resolve({
327 | text: () => 'Complete response',
328 | modelVersion: 'gemini-2.0-flash-thinking-exp',
329 | }),
330 | };
331 |
332 | const wrappedStream = meter.trackUsageStreamGemini(
333 | mockGeminiStream,
334 | 'cus_456',
335 | 'gemini-2.0-flash-thinking-exp'
336 | );
337 |
338 | for await (const _chunk of wrappedStream.stream) {
339 | // Consume stream
340 | }
341 |
342 | // Wait for fire-and-forget logging to complete
343 | await new Promise(resolve => setImmediate(resolve));
344 |
345 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
346 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
347 | expect.objectContaining({
348 | payload: expect.objectContaining({
349 | stripe_customer_id: 'cus_456',
350 | value: '20',
351 | model: 'google/gemini-2.0-flash-thinking-exp',
352 | token_type: 'input',
353 | }),
354 | })
355 | );
356 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
357 | expect.objectContaining({
358 | payload: expect.objectContaining({
359 | value: '65', // 15 candidates + 50 thoughts
360 | token_type: 'output',
361 | }),
362 | })
363 | );
364 | });
365 |
366 | it('should preserve the response promise in wrapped stream', async () => {
367 | const meter = createTokenMeter(TEST_API_KEY, config);
368 |
369 | const mockGeminiStream = {
370 | stream: {
371 | async *[Symbol.asyncIterator]() {
372 | yield {
373 | text: () => 'Hello',
374 | usageMetadata: {
375 | promptTokenCount: 10,
376 | candidatesTokenCount: 5,
377 | totalTokenCount: 15,
378 | },
379 | };
380 | },
381 | },
382 | response: Promise.resolve({
383 | text: () => 'Hello',
384 | modelVersion: 'gemini-1.5-pro',
385 | }),
386 | };
387 |
388 | const wrappedStream = meter.trackUsageStreamGemini(
389 | mockGeminiStream,
390 | 'cus_123',
391 | 'gemini-1.5-pro'
392 | );
393 |
394 | expect(wrappedStream).toHaveProperty('stream');
395 | expect(wrappedStream).toHaveProperty('response');
396 |
397 | const response = await wrappedStream.response;
398 | expect(response.text()).toBe('Hello');
399 | });
400 |
401 | it('should properly wrap the stream generator', async () => {
402 | const meter = createTokenMeter(TEST_API_KEY, config);
403 |
404 | const chunks = [
405 | {text: () => 'First', usageMetadata: null},
406 | {text: () => ' Second', usageMetadata: null},
407 | {
408 | text: () => ' Third',
409 | usageMetadata: {
410 | promptTokenCount: 20,
411 | candidatesTokenCount: 15,
412 | totalTokenCount: 35,
413 | },
414 | },
415 | ];
416 |
417 | const mockGeminiStream = {
418 | stream: {
419 | async *[Symbol.asyncIterator]() {
420 | for (const chunk of chunks) {
421 | yield chunk;
422 | }
423 | },
424 | },
425 | response: Promise.resolve({
426 | text: () => 'First Second Third',
427 | modelVersion: 'gemini-2.0-flash-exp',
428 | }),
429 | };
430 |
431 | const wrappedStream = meter.trackUsageStreamGemini(
432 | mockGeminiStream,
433 | 'cus_123',
434 | 'gemini-2.0-flash-exp'
435 | );
436 |
437 | const receivedChunks: any[] = [];
438 | for await (const chunk of wrappedStream.stream) {
439 | receivedChunks.push(chunk);
440 | }
441 |
442 | expect(receivedChunks).toHaveLength(3);
443 | expect(receivedChunks[0].text()).toBe('First');
444 | expect(receivedChunks[1].text()).toBe(' Second');
445 | expect(receivedChunks[2].text()).toBe(' Third');
446 | });
447 | });
448 |
449 | describe('Multi-turn Chat (ChatSession)', () => {
450 | it('should track usage from chat session message', async () => {
451 | const meter = createTokenMeter(TEST_API_KEY, config);
452 |
453 | // ChatSession.sendMessage() returns the same structure as generateContent
454 | const response = {
455 | response: {
456 | text: () => 'This is my second response.',
457 | usageMetadata: {
458 | promptTokenCount: 80, // Includes conversation history
459 | candidatesTokenCount: 12,
460 | totalTokenCount: 92,
461 | },
462 | modelVersion: 'gemini-2.5-flash',
463 | },
464 | };
465 |
466 | meter.trackUsage(response, 'cus_123');
467 |
468 | // Wait for fire-and-forget logging to complete
469 | await new Promise(resolve => setImmediate(resolve));
470 |
471 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
472 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
473 | expect.objectContaining({
474 | payload: expect.objectContaining({
475 | value: '80',
476 | model: 'google/gemini-2.5-flash',
477 | token_type: 'input',
478 | }),
479 | })
480 | );
481 | });
482 |
483 | it('should track usage from streaming chat session', async () => {
484 | const meter = createTokenMeter(TEST_API_KEY, config);
485 |
486 | const chunks = [
487 | {
488 | text: () => 'Continuing',
489 | usageMetadata: null,
490 | },
491 | {
492 | text: () => ' our conversation.',
493 | usageMetadata: {
494 | promptTokenCount: 100, // Includes full conversation context
495 | candidatesTokenCount: 10,
496 | totalTokenCount: 110,
497 | },
498 | },
499 | ];
500 |
501 | const mockGeminiStream = {
502 | stream: {
503 | async *[Symbol.asyncIterator]() {
504 | for (const chunk of chunks) {
505 | yield chunk;
506 | }
507 | },
508 | },
509 | response: Promise.resolve({
510 | text: () => 'Continuing our conversation.',
511 | modelVersion: 'gemini-1.5-pro',
512 | }),
513 | };
514 |
515 | const wrappedStream = meter.trackUsageStreamGemini(
516 | mockGeminiStream,
517 | 'cus_456',
518 | 'gemini-1.5-pro'
519 | );
520 |
521 | for await (const _chunk of wrappedStream.stream) {
522 | // Consume stream
523 | }
524 |
525 | // Wait for fire-and-forget logging to complete
526 | await new Promise(resolve => setImmediate(resolve));
527 |
528 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
529 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
530 | expect.objectContaining({
531 | payload: expect.objectContaining({
532 | value: '100',
533 | model: 'google/gemini-1.5-pro',
534 | token_type: 'input',
535 | }),
536 | })
537 | );
538 | });
539 |
540 | it('should track usage from long conversation with history', async () => {
541 | const meter = createTokenMeter(TEST_API_KEY, config);
542 |
543 | const response = {
544 | response: {
545 | text: () => 'Based on our previous discussion...',
546 | usageMetadata: {
547 | promptTokenCount: 500, // Large context from history
548 | candidatesTokenCount: 25,
549 | totalTokenCount: 525,
550 | },
551 | modelVersion: 'gemini-1.5-pro',
552 | },
553 | };
554 |
555 | meter.trackUsage(response, 'cus_789');
556 |
557 | // Wait for fire-and-forget logging to complete
558 | await new Promise(resolve => setImmediate(resolve));
559 |
560 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
561 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
562 | expect.objectContaining({
563 | payload: expect.objectContaining({
564 | value: '500',
565 | model: 'google/gemini-1.5-pro',
566 | token_type: 'input',
567 | }),
568 | })
569 | );
570 | });
571 | });
572 |
573 | describe('Model Variants', () => {
574 | it('should track gemini-1.5-pro', async () => {
575 | const meter = createTokenMeter(TEST_API_KEY, config);
576 |
577 | const response = {
578 | response: {
579 | text: () => 'Pro model response',
580 | usageMetadata: {
581 | promptTokenCount: 15,
582 | candidatesTokenCount: 10,
583 | totalTokenCount: 25,
584 | },
585 | modelVersion: 'gemini-1.5-pro',
586 | },
587 | };
588 |
589 | meter.trackUsage(response, 'cus_123');
590 |
591 | // Wait for fire-and-forget logging to complete
592 | await new Promise(resolve => setImmediate(resolve));
593 |
594 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
595 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
596 | expect.objectContaining({
597 | payload: expect.objectContaining({
598 | value: '15',
599 | model: 'google/gemini-1.5-pro',
600 | token_type: 'input',
601 | }),
602 | })
603 | );
604 | });
605 |
606 | it('should track gemini-2.5-flash', async () => {
607 | const meter = createTokenMeter(TEST_API_KEY, config);
608 |
609 | const response = {
610 | response: {
611 | text: () => 'Flash model response',
612 | usageMetadata: {
613 | promptTokenCount: 12,
614 | candidatesTokenCount: 8,
615 | totalTokenCount: 20,
616 | },
617 | modelVersion: 'gemini-2.5-flash',
618 | },
619 | };
620 |
621 | meter.trackUsage(response, 'cus_456');
622 |
623 | // Wait for fire-and-forget logging to complete
624 | await new Promise(resolve => setImmediate(resolve));
625 |
626 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
627 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
628 | expect.objectContaining({
629 | payload: expect.objectContaining({
630 | value: '12',
631 | model: 'google/gemini-2.5-flash',
632 | token_type: 'input',
633 | }),
634 | })
635 | );
636 | });
637 |
638 | it('should track gemini-2.0-flash-exp', async () => {
639 | const meter = createTokenMeter(TEST_API_KEY, config);
640 |
641 | const response = {
642 | response: {
643 | text: () => 'Gemini 2.0 response',
644 | usageMetadata: {
645 | promptTokenCount: 10,
646 | candidatesTokenCount: 5,
647 | totalTokenCount: 15,
648 | },
649 | modelVersion: 'gemini-2.0-flash-exp',
650 | },
651 | };
652 |
653 | meter.trackUsage(response, 'cus_789');
654 |
655 | // Wait for fire-and-forget logging to complete
656 | await new Promise(resolve => setImmediate(resolve));
657 |
658 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
659 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
660 | expect.objectContaining({
661 | payload: expect.objectContaining({
662 | value: '10',
663 | model: 'google/gemini-2.0-flash-exp',
664 | token_type: 'input',
665 | }),
666 | })
667 | );
668 | });
669 |
670 | it('should track gemini-2.0-flash-thinking-exp with reasoning tokens', async () => {
671 | const meter = createTokenMeter(TEST_API_KEY, config);
672 |
673 | const response = {
674 | response: {
675 | text: () => 'Thought-through response',
676 | usageMetadata: {
677 | promptTokenCount: 25,
678 | candidatesTokenCount: 20,
679 | thoughtsTokenCount: 100, // Extended thinking
680 | totalTokenCount: 145,
681 | },
682 | modelVersion: 'gemini-2.0-flash-thinking-exp',
683 | },
684 | };
685 |
686 | meter.trackUsage(response, 'cus_999');
687 |
688 | // Wait for fire-and-forget logging to complete
689 | await new Promise(resolve => setImmediate(resolve));
690 |
691 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
692 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
693 | expect.objectContaining({
694 | payload: expect.objectContaining({
695 | value: '25',
696 | model: 'google/gemini-2.0-flash-thinking-exp',
697 | token_type: 'input',
698 | }),
699 | })
700 | );
701 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
702 | expect.objectContaining({
703 | payload: expect.objectContaining({
704 | value: '120', // 20 + 100
705 | token_type: 'output',
706 | }),
707 | })
708 | );
709 | });
710 | });
711 |
712 | describe('Edge Cases', () => {
713 | it('should handle zero reasoning tokens gracefully', async () => {
714 | const meter = createTokenMeter(TEST_API_KEY, config);
715 |
716 | const response = {
717 | response: {
718 | text: () => 'No reasoning tokens',
719 | usageMetadata: {
720 | promptTokenCount: 10,
721 | candidatesTokenCount: 5,
722 | thoughtsTokenCount: 0,
723 | totalTokenCount: 15,
724 | },
725 | modelVersion: 'gemini-2.0-flash-thinking-exp',
726 | },
727 | };
728 |
729 | meter.trackUsage(response, 'cus_123');
730 |
731 | // Wait for fire-and-forget logging to complete
732 | await new Promise(resolve => setImmediate(resolve));
733 |
734 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
735 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
736 | expect.objectContaining({
737 | payload: expect.objectContaining({
738 | value: '10',
739 | model: 'google/gemini-2.0-flash-thinking-exp',
740 | token_type: 'input',
741 | }),
742 | })
743 | );
744 | });
745 |
746 | it('should handle missing thoughtsTokenCount field', async () => {
747 | const meter = createTokenMeter(TEST_API_KEY, config);
748 |
749 | const response = {
750 | response: {
751 | text: () => 'Standard model without thoughts',
752 | usageMetadata: {
753 | promptTokenCount: 10,
754 | candidatesTokenCount: 5,
755 | totalTokenCount: 15,
756 | // No thoughtsTokenCount field
757 | },
758 | modelVersion: 'gemini-1.5-pro',
759 | },
760 | };
761 |
762 | meter.trackUsage(response, 'cus_123');
763 |
764 | // Wait for fire-and-forget logging to complete
765 | await new Promise(resolve => setImmediate(resolve));
766 |
767 | expect(mockMeterEventsCreate).toHaveBeenCalledTimes(2);
768 | expect(mockMeterEventsCreate).toHaveBeenCalledWith(
769 | expect.objectContaining({
770 | payload: expect.objectContaining({
771 | value: '10',
772 | model: 'google/gemini-1.5-pro',
773 | token_type: 'input',
774 | }),
775 | })
776 | );
777 | });
778 | });
779 | });
780 |
781 |
```
--------------------------------------------------------------------------------
/tools/python/tests/test_functions.py:
--------------------------------------------------------------------------------
```python
1 | import unittest
2 | import stripe
3 | from unittest import mock
4 | from stripe_agent_toolkit.functions import (
5 | create_customer,
6 | list_customers,
7 | create_product,
8 | list_products,
9 | create_price,
10 | list_prices,
11 | create_payment_link,
12 | list_invoices,
13 | create_invoice,
14 | create_invoice_item,
15 | finalize_invoice,
16 | retrieve_balance,
17 | create_refund,
18 | list_payment_intents,
19 | create_billing_portal_session,
20 | )
21 |
22 |
23 | class TestStripeFunctions(unittest.TestCase):
24 | def test_create_customer(self):
25 | with mock.patch("stripe.Customer.create") as mock_function:
26 | mock_customer = {"id": "cus_123"}
27 | mock_function.return_value = stripe.Customer.construct_from(
28 | mock_customer, "sk_test_123"
29 | )
30 |
31 | result = create_customer(
32 | context={}, name="Test User", email="[email protected]"
33 | )
34 |
35 | mock_function.assert_called_with(
36 | name="Test User", email="[email protected]"
37 | )
38 |
39 | self.assertEqual(result, {"id": mock_customer["id"]})
40 |
41 | def test_create_customer_with_context(self):
42 | with mock.patch("stripe.Customer.create") as mock_function:
43 | mock_customer = {"id": "cus_123"}
44 | mock_function.return_value = stripe.Customer.construct_from(
45 | mock_customer, "sk_test_123"
46 | )
47 |
48 | result = create_customer(
49 | context={"account": "acct_123"},
50 | name="Test User",
51 | email="[email protected]",
52 | )
53 |
54 | mock_function.assert_called_with(
55 | name="Test User",
56 | email="[email protected]",
57 | stripe_account="acct_123",
58 | )
59 |
60 | self.assertEqual(result, {"id": mock_customer["id"]})
61 |
62 | def test_list_customers(self):
63 | with mock.patch("stripe.Customer.list") as mock_function:
64 | mock_customers = [{"id": "cus_123"}, {"id": "cus_456"}]
65 |
66 | mock_function.return_value = stripe.ListObject.construct_from(
67 | {
68 | "object": "list",
69 | "data": [
70 | stripe.Customer.construct_from(
71 | {
72 | "id": "cus_123",
73 | "email": "[email protected]",
74 | "name": "Customer One",
75 | },
76 | "sk_test_123",
77 | ),
78 | stripe.Customer.construct_from(
79 | {
80 | "id": "cus_456",
81 | "email": "[email protected]",
82 | "name": "Customer Two",
83 | },
84 | "sk_test_123",
85 | ),
86 | ],
87 | "has_more": False,
88 | "url": "/v1/customers",
89 | },
90 | "sk_test_123",
91 | )
92 |
93 | result = list_customers(context={})
94 |
95 | mock_function.assert_called_with()
96 |
97 | self.assertEqual(result, mock_customers)
98 |
99 | def test_list_customers_with_context(self):
100 | with mock.patch("stripe.Customer.list") as mock_function:
101 | mock_customers = [{"id": "cus_123"}, {"id": "cus_456"}]
102 |
103 | mock_function.return_value = stripe.ListObject.construct_from(
104 | {
105 | "object": "list",
106 | "data": [
107 | stripe.Customer.construct_from(
108 | {
109 | "id": "cus_123",
110 | "email": "[email protected]",
111 | "name": "Customer One",
112 | },
113 | "sk_test_123",
114 | ),
115 | stripe.Customer.construct_from(
116 | {
117 | "id": "cus_456",
118 | "email": "[email protected]",
119 | "name": "Customer Two",
120 | },
121 | "sk_test_123",
122 | ),
123 | ],
124 | "has_more": False,
125 | "url": "/v1/customers",
126 | },
127 | "sk_test_123",
128 | )
129 |
130 | result = list_customers(context={"account": "acct_123"})
131 |
132 | mock_function.assert_called_with(
133 | stripe_account="acct_123",
134 | )
135 |
136 | self.assertEqual(result, mock_customers)
137 |
138 | def test_create_product(self):
139 | with mock.patch("stripe.Product.create") as mock_function:
140 | mock_product = {"id": "prod_123"}
141 | mock_function.return_value = stripe.Product.construct_from(
142 | mock_product, "sk_test_123"
143 | )
144 |
145 | result = create_product(context={}, name="Test Product")
146 |
147 | mock_function.assert_called_with(
148 | name="Test Product",
149 | )
150 |
151 | self.assertEqual(result, {"id": mock_product["id"]})
152 |
153 | def test_create_product_with_context(self):
154 | with mock.patch("stripe.Product.create") as mock_function:
155 | mock_product = {"id": "prod_123"}
156 | mock_function.return_value = stripe.Product.construct_from(
157 | mock_product, "sk_test_123"
158 | )
159 |
160 | result = create_product(
161 | context={"account": "acct_123"}, name="Test Product"
162 | )
163 |
164 | mock_function.assert_called_with(
165 | name="Test Product", stripe_account="acct_123"
166 | )
167 |
168 | self.assertEqual(result, {"id": mock_product["id"]})
169 |
170 | def test_list_products(self):
171 | with mock.patch("stripe.Product.list") as mock_function:
172 | mock_products = [
173 | {"id": "prod_123", "name": "Product One"},
174 | {"id": "prod_456", "name": "Product Two"},
175 | ]
176 |
177 | mock_function.return_value = stripe.ListObject.construct_from(
178 | {
179 | "object": "list",
180 | "data": [
181 | stripe.Product.construct_from(
182 | {
183 | "id": "prod_123",
184 | "name": "Product One",
185 | },
186 | "sk_test_123",
187 | ),
188 | stripe.Product.construct_from(
189 | {
190 | "id": "prod_456",
191 | "name": "Product Two",
192 | },
193 | "sk_test_123",
194 | ),
195 | ],
196 | "has_more": False,
197 | "url": "/v1/products",
198 | },
199 | "sk_test_123",
200 | )
201 |
202 | result = list_products(context={})
203 |
204 | mock_function.assert_called_with()
205 |
206 | self.assertEqual(result, mock_products)
207 |
208 | def test_create_price(self):
209 | with mock.patch("stripe.Price.create") as mock_function:
210 | mock_price = {"id": "price_123"}
211 | mock_function.return_value = stripe.Price.construct_from(
212 | mock_price, "sk_test_123"
213 | )
214 |
215 | result = create_price(
216 | context={},
217 | product="prod_123",
218 | currency="usd",
219 | unit_amount=1000,
220 | )
221 |
222 | mock_function.assert_called_with(
223 | product="prod_123",
224 | currency="usd",
225 | unit_amount=1000,
226 | )
227 |
228 | self.assertEqual(result, {"id": mock_price["id"]})
229 |
230 | def test_create_price_with_context(self):
231 | with mock.patch("stripe.Price.create") as mock_function:
232 | mock_price = {"id": "price_123"}
233 | mock_function.return_value = stripe.Price.construct_from(
234 | mock_price, "sk_test_123"
235 | )
236 |
237 | result = create_price(
238 | context={"account": "acct_123"},
239 | product="prod_123",
240 | currency="usd",
241 | unit_amount=1000,
242 | )
243 |
244 | mock_function.assert_called_with(
245 | product="prod_123",
246 | currency="usd",
247 | unit_amount=1000,
248 | stripe_account="acct_123",
249 | )
250 |
251 | self.assertEqual(result, {"id": mock_price["id"]})
252 |
253 | def test_list_prices(self):
254 | with mock.patch("stripe.Price.list") as mock_function:
255 | mock_prices = [
256 | {"id": "price_123", "product": "prod_123"},
257 | {"id": "price_456", "product": "prod_456"},
258 | ]
259 |
260 | mock_function.return_value = stripe.ListObject.construct_from(
261 | {
262 | "object": "list",
263 | "data": [
264 | stripe.Price.construct_from(
265 | {
266 | "id": "price_123",
267 | "product": "prod_123",
268 | },
269 | "sk_test_123",
270 | ),
271 | stripe.Price.construct_from(
272 | {
273 | "id": "price_456",
274 | "product": "prod_456",
275 | },
276 | "sk_test_123",
277 | ),
278 | ],
279 | "has_more": False,
280 | "url": "/v1/prices",
281 | },
282 | "sk_test_123",
283 | )
284 |
285 | result = list_prices({})
286 |
287 | mock_function.assert_called_with()
288 |
289 | self.assertEqual(result, mock_prices)
290 |
291 | def test_list_prices_with_context(self):
292 | with mock.patch("stripe.Price.list") as mock_function:
293 | mock_prices = [
294 | {"id": "price_123", "product": "prod_123"},
295 | {"id": "price_456", "product": "prod_456"},
296 | ]
297 |
298 | mock_function.return_value = stripe.ListObject.construct_from(
299 | {
300 | "object": "list",
301 | "data": [
302 | stripe.Price.construct_from(
303 | {
304 | "id": "price_123",
305 | "product": "prod_123",
306 | },
307 | "sk_test_123",
308 | ),
309 | stripe.Price.construct_from(
310 | {
311 | "id": "price_456",
312 | "product": "prod_456",
313 | },
314 | "sk_test_123",
315 | ),
316 | ],
317 | "has_more": False,
318 | "url": "/v1/prices",
319 | },
320 | "sk_test_123",
321 | )
322 |
323 | result = list_prices({"account": "acct_123"})
324 |
325 | mock_function.assert_called_with(stripe_account="acct_123")
326 |
327 | self.assertEqual(result, mock_prices)
328 |
329 | def test_create_payment_link(self):
330 | with mock.patch("stripe.PaymentLink.create") as mock_function:
331 | mock_payment_link = {"id": "pl_123", "url": "https://example.com"}
332 | mock_function.return_value = stripe.PaymentLink.construct_from(
333 | mock_payment_link, "sk_test_123"
334 | )
335 |
336 | result = create_payment_link(
337 | context={}, price="price_123", quantity=1
338 | )
339 |
340 | mock_function.assert_called_with(
341 | line_items=[{"price": "price_123", "quantity": 1}],
342 | )
343 |
344 | self.assertEqual(result, mock_payment_link)
345 |
346 | def test_create_payment_link_with_redirect_url(self):
347 | with mock.patch("stripe.PaymentLink.create") as mock_function:
348 | mock_payment_link = {"id": "pl_123", "url": "https://example.com"}
349 | mock_function.return_value = stripe.PaymentLink.construct_from(
350 | mock_payment_link, "sk_test_123"
351 | )
352 |
353 | result = create_payment_link(
354 | context={}, price="price_123", quantity=1, redirect_url="https://example.com"
355 | )
356 |
357 | mock_function.assert_called_with(
358 | line_items=[{"price": "price_123", "quantity": 1, }],
359 | after_completion={"type": "redirect", "redirect": {"url": "https://example.com"}}
360 | )
361 |
362 | self.assertEqual(result, mock_payment_link)
363 |
364 | def test_create_payment_link_with_context(self):
365 | with mock.patch("stripe.PaymentLink.create") as mock_function:
366 | mock_payment_link = {"id": "pl_123", "url": "https://example.com"}
367 | mock_function.return_value = stripe.PaymentLink.construct_from(
368 | mock_payment_link, "sk_test_123"
369 | )
370 |
371 | result = create_payment_link(
372 | context={"account": "acct_123"}, price="price_123", quantity=1
373 | )
374 |
375 | mock_function.assert_called_with(
376 | line_items=[{"price": "price_123", "quantity": 1}],
377 | stripe_account="acct_123",
378 | )
379 |
380 | self.assertEqual(result, mock_payment_link)
381 |
382 | def test_list_invoices(self):
383 | with mock.patch("stripe.Invoice.list") as mock_function:
384 | mock_invoice = {
385 | "id": "in_123",
386 | "hosted_invoice_url": "https://example.com",
387 | "customer": "cus_123",
388 | "status": "open",
389 | }
390 | mock_invoices = {
391 | "object": "list",
392 | "data": [
393 | stripe.Invoice.construct_from(
394 | mock_invoice,
395 | "sk_test_123",
396 | ),
397 | ],
398 | "has_more": False,
399 | "url": "/v1/invoices",
400 | }
401 |
402 | mock_function.return_value = stripe.Invoice.construct_from(
403 | mock_invoices, "sk_test_123"
404 | )
405 |
406 | result = list_invoices(context={})
407 |
408 | mock_function.assert_called_with()
409 |
410 | self.assertEqual(
411 | result,
412 | [
413 | {
414 | "id": mock_invoice["id"],
415 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
416 | "customer": mock_invoice["customer"],
417 | "status": mock_invoice["status"],
418 | }
419 | ],
420 | )
421 |
422 | def test_list_invoices_with_customer(self):
423 | with mock.patch("stripe.Invoice.list") as mock_function:
424 | mock_invoice = {
425 | "id": "in_123",
426 | "hosted_invoice_url": "https://example.com",
427 | "customer": "cus_123",
428 | "status": "open",
429 | }
430 | mock_invoices = {
431 | "object": "list",
432 | "data": [
433 | stripe.Invoice.construct_from(
434 | mock_invoice,
435 | "sk_test_123",
436 | ),
437 | ],
438 | "has_more": False,
439 | "url": "/v1/invoices",
440 | }
441 |
442 | mock_function.return_value = stripe.Invoice.construct_from(
443 | mock_invoices, "sk_test_123"
444 | )
445 |
446 | result = list_invoices(context={}, customer="cus_123")
447 |
448 | mock_function.assert_called_with(
449 | customer="cus_123",
450 | )
451 |
452 | self.assertEqual(
453 | result,
454 | [
455 | {
456 | "id": mock_invoice["id"],
457 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
458 | "customer": mock_invoice["customer"],
459 | "status": mock_invoice["status"],
460 | }
461 | ],
462 | )
463 |
464 | def test_list_invoices_with_customer_and_limit(self):
465 | with mock.patch("stripe.Invoice.list") as mock_function:
466 | mock_invoice = {
467 | "id": "in_123",
468 | "hosted_invoice_url": "https://example.com",
469 | "customer": "cus_123",
470 | "status": "open",
471 | }
472 | mock_invoices = {
473 | "object": "list",
474 | "data": [
475 | stripe.Invoice.construct_from(
476 | mock_invoice,
477 | "sk_test_123",
478 | ),
479 | ],
480 | "has_more": False,
481 | "url": "/v1/invoices",
482 | }
483 |
484 | mock_function.return_value = stripe.Invoice.construct_from(
485 | mock_invoices, "sk_test_123"
486 | )
487 |
488 | result = list_invoices(context={}, customer="cus_123", limit=100)
489 |
490 | mock_function.assert_called_with(
491 | customer="cus_123",
492 | limit=100,
493 | )
494 |
495 | self.assertEqual(
496 | result,
497 | [
498 | {
499 | "id": mock_invoice["id"],
500 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
501 | "customer": mock_invoice["customer"],
502 | "status": mock_invoice["status"],
503 | }
504 | ],
505 | )
506 |
507 | def test_list_invoices_with_context(self):
508 | with mock.patch("stripe.Invoice.list") as mock_function:
509 | mock_invoice = {
510 | "id": "in_123",
511 | "hosted_invoice_url": "https://example.com",
512 | "customer": "cus_123",
513 | "status": "open",
514 | }
515 | mock_invoices = {
516 | "object": "list",
517 | "data": [
518 | stripe.Invoice.construct_from(
519 | mock_invoice,
520 | "sk_test_123",
521 | ),
522 | ],
523 | "has_more": False,
524 | "url": "/v1/invoices",
525 | }
526 |
527 | mock_function.return_value = stripe.Invoice.construct_from(
528 | mock_invoices, "sk_test_123"
529 | )
530 |
531 | result = list_invoices(context={"account": "acct_123"}, customer="cus_123")
532 |
533 | mock_function.assert_called_with(
534 | customer="cus_123",
535 | stripe_account="acct_123",
536 | )
537 |
538 | self.assertEqual(
539 | result,
540 | [
541 | {
542 | "id": mock_invoice["id"],
543 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
544 | "customer": mock_invoice["customer"],
545 | "status": mock_invoice["status"],
546 | }
547 | ],
548 | )
549 |
550 | def test_create_invoice(self):
551 | with mock.patch("stripe.Invoice.create") as mock_function:
552 | mock_invoice = {
553 | "id": "in_123",
554 | "hosted_invoice_url": "https://example.com",
555 | "customer": "cus_123",
556 | "status": "open",
557 | }
558 |
559 | mock_function.return_value = stripe.Invoice.construct_from(
560 | mock_invoice, "sk_test_123"
561 | )
562 |
563 | result = create_invoice(context={}, customer="cus_123")
564 |
565 | mock_function.assert_called_with(
566 | customer="cus_123",
567 | collection_method="send_invoice",
568 | days_until_due=30,
569 | )
570 |
571 | self.assertEqual(
572 | result,
573 | {
574 | "id": mock_invoice["id"],
575 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
576 | "customer": mock_invoice["customer"],
577 | "status": mock_invoice["status"],
578 | },
579 | )
580 |
581 | def test_create_invoice_with_context(self):
582 | with mock.patch("stripe.Invoice.create") as mock_function:
583 | mock_invoice = {
584 | "id": "in_123",
585 | "hosted_invoice_url": "https://example.com",
586 | "customer": "cus_123",
587 | "status": "open",
588 | }
589 |
590 | mock_function.return_value = stripe.Invoice.construct_from(
591 | mock_invoice, "sk_test_123"
592 | )
593 |
594 | result = create_invoice(
595 | context={"account": "acct_123"}, customer="cus_123"
596 | )
597 |
598 | mock_function.assert_called_with(
599 | customer="cus_123",
600 | collection_method="send_invoice",
601 | days_until_due=30,
602 | stripe_account="acct_123",
603 | )
604 |
605 | self.assertEqual(
606 | result,
607 | {
608 | "id": mock_invoice["id"],
609 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
610 | "customer": mock_invoice["customer"],
611 | "status": mock_invoice["status"],
612 | },
613 | )
614 |
615 | def test_create_invoice_item(self):
616 | with mock.patch("stripe.InvoiceItem.create") as mock_function:
617 | mock_invoice_item = {"id": "ii_123", "invoice": "in_123"}
618 | mock_function.return_value = stripe.InvoiceItem.construct_from(
619 | mock_invoice_item, "sk_test_123"
620 | )
621 |
622 | result = create_invoice_item(
623 | context={},
624 | customer="cus_123",
625 | price="price_123",
626 | invoice="in_123",
627 | )
628 |
629 | mock_function.assert_called_with(
630 | customer="cus_123", price="price_123", invoice="in_123"
631 | )
632 |
633 | self.assertEqual(
634 | result,
635 | {
636 | "id": mock_invoice_item["id"],
637 | "invoice": mock_invoice_item["invoice"],
638 | },
639 | )
640 |
641 | def test_create_invoice_item_with_context(self):
642 | with mock.patch("stripe.InvoiceItem.create") as mock_function:
643 | mock_invoice_item = {"id": "ii_123", "invoice": "in_123"}
644 | mock_function.return_value = stripe.InvoiceItem.construct_from(
645 | mock_invoice_item, "sk_test_123"
646 | )
647 |
648 | result = create_invoice_item(
649 | context={"account": "acct_123"},
650 | customer="cus_123",
651 | price="price_123",
652 | invoice="in_123",
653 | )
654 |
655 | mock_function.assert_called_with(
656 | customer="cus_123",
657 | price="price_123",
658 | invoice="in_123",
659 | stripe_account="acct_123",
660 | )
661 |
662 | self.assertEqual(
663 | result,
664 | {
665 | "id": mock_invoice_item["id"],
666 | "invoice": mock_invoice_item["invoice"],
667 | },
668 | )
669 |
670 | def test_finalize_invoice(self):
671 | with mock.patch("stripe.Invoice.finalize_invoice") as mock_function:
672 | mock_invoice = {
673 | "id": "in_123",
674 | "hosted_invoice_url": "https://example.com",
675 | "customer": "cus_123",
676 | "status": "open",
677 | }
678 |
679 | mock_function.return_value = stripe.Invoice.construct_from(
680 | mock_invoice, "sk_test_123"
681 | )
682 |
683 | result = finalize_invoice(context={}, invoice="in_123")
684 |
685 | mock_function.assert_called_with(invoice="in_123")
686 |
687 | self.assertEqual(
688 | result,
689 | {
690 | "id": mock_invoice["id"],
691 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
692 | "customer": mock_invoice["customer"],
693 | "status": mock_invoice["status"],
694 | },
695 | )
696 |
697 | def test_finalize_invoice_with_context(self):
698 | with mock.patch("stripe.Invoice.finalize_invoice") as mock_function:
699 | mock_invoice = {
700 | "id": "in_123",
701 | "hosted_invoice_url": "https://example.com",
702 | "customer": "cus_123",
703 | "status": "open",
704 | }
705 |
706 | mock_function.return_value = stripe.Invoice.construct_from(
707 | mock_invoice, "sk_test_123"
708 | )
709 |
710 | result = finalize_invoice(
711 | context={"account": "acct_123"}, invoice="in_123"
712 | )
713 |
714 | mock_function.assert_called_with(
715 | invoice="in_123", stripe_account="acct_123"
716 | )
717 |
718 | self.assertEqual(
719 | result,
720 | {
721 | "id": mock_invoice["id"],
722 | "hosted_invoice_url": mock_invoice["hosted_invoice_url"],
723 | "customer": mock_invoice["customer"],
724 | "status": mock_invoice["status"],
725 | },
726 | )
727 |
728 | def test_retrieve_balance(self):
729 | with mock.patch("stripe.Balance.retrieve") as mock_function:
730 | mock_balance = {"available": [{"amount": 1000, "currency": "usd"}]}
731 |
732 | mock_function.return_value = stripe.Balance.construct_from(
733 | mock_balance, "sk_test_123"
734 | )
735 |
736 | result = retrieve_balance(context={})
737 |
738 | mock_function.assert_called_with()
739 |
740 | self.assertEqual(result, mock_balance)
741 |
742 | def test_retrieve_balance_with_context(self):
743 | with mock.patch("stripe.Balance.retrieve") as mock_function:
744 | mock_balance = {"available": [{"amount": 1000, "currency": "usd"}]}
745 |
746 | mock_function.return_value = stripe.Balance.construct_from(
747 | mock_balance, "sk_test_123"
748 | )
749 |
750 | result = retrieve_balance(context={"account": "acct_123"})
751 |
752 | mock_function.assert_called_with(stripe_account="acct_123")
753 |
754 | self.assertEqual(result, mock_balance)
755 |
756 | def test_create_refund(self):
757 | with mock.patch("stripe.Refund.create") as mock_function:
758 | mock_refund = {"id": "re_123"}
759 | mock_function.return_value = stripe.Refund.construct_from(
760 | mock_refund, "sk_test_123"
761 | )
762 |
763 | result = create_refund(context={}, payment_intent="pi_123")
764 |
765 | mock_function.assert_called_with(payment_intent="pi_123")
766 |
767 | self.assertEqual(result, {"id": mock_refund["id"]})
768 |
769 | def test_create_partial_refund(self):
770 | with mock.patch("stripe.Refund.create") as mock_function:
771 | mock_refund = {"id": "re_123"}
772 | mock_function.return_value = stripe.Refund.construct_from(
773 | mock_refund, "sk_test_123"
774 | )
775 |
776 | result = create_refund(
777 | context={}, payment_intent="pi_123", amount=1000
778 | )
779 |
780 | mock_function.assert_called_with(
781 | payment_intent="pi_123", amount=1000
782 | )
783 |
784 | self.assertEqual(result, {"id": mock_refund["id"]})
785 |
786 | def test_create_refund_with_context(self):
787 | with mock.patch("stripe.Refund.create") as mock_function:
788 | mock_refund = {"id": "re_123"}
789 | mock_function.return_value = stripe.Refund.construct_from(
790 | mock_refund, "sk_test_123"
791 | )
792 |
793 | result = create_refund(
794 | context={"account": "acct_123"},
795 | payment_intent="pi_123",
796 | amount=1000,
797 | )
798 |
799 | mock_function.assert_called_with(
800 | payment_intent="pi_123", amount=1000, stripe_account="acct_123"
801 | )
802 |
803 | self.assertEqual(result, {"id": mock_refund["id"]})
804 |
805 | def test_list_payment_intents(self):
806 | with mock.patch("stripe.PaymentIntent.list") as mock_function:
807 | mock_payment_intents = [{"id": "pi_123"}, {"id": "pi_456"}]
808 | mock_function.return_value = stripe.ListObject.construct_from(
809 | {"data": mock_payment_intents}, "sk_test_123"
810 | )
811 |
812 | result = list_payment_intents(context={})
813 |
814 | mock_function.assert_called_with()
815 |
816 | self.assertEqual(result, mock_payment_intents)
817 |
818 | def test_list_payment_intents_with_context(self):
819 | with mock.patch("stripe.PaymentIntent.list") as mock_function:
820 | mock_payment_intents = [{"id": "pi_123"}, {"id": "pi_456"}]
821 | mock_function.return_value = stripe.ListObject.construct_from(
822 | {"data": mock_payment_intents}, "sk_test_123"
823 | )
824 |
825 | result = list_payment_intents(context={"account": "acct_123"})
826 |
827 | mock_function.assert_called_with(stripe_account="acct_123")
828 |
829 | self.assertEqual(result, mock_payment_intents)
830 |
831 |
832 | def test_create_billing_portal_session(self):
833 | with mock.patch("stripe.billing_portal.Session.create") as mock_function:
834 | mock_billing_portal_session = {
835 | "id": "bps_123",
836 | "url": "https://example.com",
837 | "customer": "cus_123",
838 | "configuration": "bpc_123",
839 | }
840 | mock_function.return_value = stripe.billing_portal.Session.construct_from(
841 | mock_billing_portal_session, "sk_test_123"
842 | )
843 |
844 | result = create_billing_portal_session(context={}, customer="cus_123")
845 |
846 | mock_function.assert_called_with(customer="cus_123")
847 |
848 | self.assertEqual(result, {
849 | "id": mock_billing_portal_session["id"],
850 | "url": mock_billing_portal_session["url"],
851 | "customer": mock_billing_portal_session["customer"],
852 | })
853 |
854 | def test_create_billing_portal_session_with_return_url(self):
855 | with mock.patch("stripe.billing_portal.Session.create") as mock_function:
856 | mock_billing_portal_session = {
857 | "id": "bps_123",
858 | "url": "https://example.com",
859 | "customer": "cus_123",
860 | "configuration": "bpc_123",
861 | }
862 | mock_function.return_value = stripe.billing_portal.Session.construct_from(
863 | mock_billing_portal_session, "sk_test_123"
864 | )
865 |
866 | result = create_billing_portal_session(
867 | context={},
868 | customer="cus_123",
869 | return_url="http://example.com"
870 | )
871 |
872 | mock_function.assert_called_with(
873 | customer="cus_123",
874 | return_url="http://example.com",
875 | )
876 |
877 | self.assertEqual(result, {
878 | "id": mock_billing_portal_session["id"],
879 | "url": mock_billing_portal_session["url"],
880 | "customer": mock_billing_portal_session["customer"],
881 | })
882 |
883 | def test_create_billing_portal_session_with_context(self):
884 | with mock.patch("stripe.billing_portal.Session.create") as mock_function:
885 | mock_billing_portal_session = {
886 | "id": "bps_123",
887 | "url": "https://example.com",
888 | "customer": "cus_123",
889 | "configuration": "bpc_123",
890 | }
891 | mock_function.return_value = stripe.billing_portal.Session.construct_from(
892 | mock_billing_portal_session, "sk_test_123"
893 | )
894 |
895 | result = create_billing_portal_session(
896 | context={"account": "acct_123"},
897 | customer="cus_123",
898 | return_url="http://example.com"
899 | )
900 |
901 | mock_function.assert_called_with(
902 | customer="cus_123",
903 | return_url="http://example.com",
904 | stripe_account="acct_123"
905 | )
906 |
907 | self.assertEqual(result, {
908 | "id": mock_billing_portal_session["id"],
909 | "url": mock_billing_portal_session["url"],
910 | "customer": mock_billing_portal_session["customer"],
911 | })
912 |
913 | if __name__ == "__main__":
914 | unittest.main()
915 |
```