Skip to content

Commit d2d42b2

Browse files
feat(billing): add generic BillingUsage model for quota tracking
1 parent 0eede6e commit d2d42b2

5 files changed

Lines changed: 311 additions & 0 deletions

File tree

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
/**
2+
* Module dependencies
3+
*/
4+
import mongoose from 'mongoose';
5+
6+
const Schema = mongoose.Schema;
7+
8+
/**
9+
* Data Model Mongoose
10+
*/
11+
const UsageMongoose = new Schema(
12+
{
13+
organizationId: {
14+
type: Schema.ObjectId,
15+
ref: 'Organization',
16+
required: true,
17+
index: true,
18+
},
19+
month: {
20+
type: String,
21+
required: true,
22+
index: true,
23+
},
24+
counters: {
25+
type: Schema.Types.Mixed,
26+
default: () => ({}),
27+
},
28+
},
29+
{
30+
timestamps: true,
31+
},
32+
);
33+
34+
UsageMongoose.index({ organizationId: 1, month: 1 }, { unique: true });
35+
36+
/**
37+
* Returns the hex string representation of the document ObjectId.
38+
* @returns {string} Hex string of the ObjectId.
39+
*/
40+
function addID() {
41+
return this._id.toHexString();
42+
}
43+
44+
/**
45+
* Model configuration
46+
*/
47+
UsageMongoose.virtual('id').get(addID);
48+
// Ensure virtual fields are serialised.
49+
UsageMongoose.set('toJSON', {
50+
virtuals: true,
51+
});
52+
53+
mongoose.model('BillingUsage', UsageMongoose);
Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
/**
2+
* Module dependencies
3+
*/
4+
import { z } from 'zod';
5+
6+
/**
7+
* Data Schema
8+
*/
9+
const objectIdRegex = /^[a-f\d]{24}$/i;
10+
11+
const BillingUsage = z.object({
12+
organizationId: z.string().trim().regex(objectIdRegex, 'organizationId must be a valid ObjectId'),
13+
month: z.string().trim().regex(/^\d{4}-\d{2}$/, 'month must be in YYYY-MM format'),
14+
counters: z.record(z.string(), z.number()).default(() => ({})),
15+
});
16+
17+
export default {
18+
BillingUsage,
19+
};
Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
1+
/**
2+
* Module dependencies
3+
*/
4+
import mongoose from 'mongoose';
5+
6+
const BillingUsage = mongoose.model('BillingUsage');
7+
8+
/**
9+
* @function get
10+
* @description Fetch a single usage document by organizationId and month.
11+
* @param {String} organizationId - The organization ID.
12+
* @param {String} month - The month in YYYY-MM format.
13+
* @returns {Promise<Object|null>} The usage document or null.
14+
*/
15+
const get = (organizationId, month) => {
16+
if (!mongoose.Types.ObjectId.isValid(organizationId)) return null;
17+
return BillingUsage.findOne({ organizationId, month }).exec();
18+
};
19+
20+
/**
21+
* @function increment
22+
* @description Atomically increment a counter key for the given org+month, with upsert.
23+
* @param {String} organizationId - The organization ID.
24+
* @param {String} month - The month in YYYY-MM format.
25+
* @param {String} key - The counter key to increment (e.g. 'executions').
26+
* @param {Number} amount - The amount to increment by.
27+
* @returns {Promise<Object>} The updated usage document.
28+
*/
29+
const increment = (organizationId, month, key, amount) => BillingUsage.findOneAndUpdate(
30+
{ organizationId, month },
31+
{ $inc: { [`counters.${key}`]: amount } },
32+
{ upsert: true, returnDocument: 'after', runValidators: true },
33+
).exec();
34+
35+
/**
36+
* @function reset
37+
* @description Reset all counters to an empty object for the given org+month.
38+
* @param {String} organizationId - The organization ID.
39+
* @param {String} month - The month in YYYY-MM format.
40+
* @returns {Promise<Object|null>} The updated usage document or null.
41+
*/
42+
const reset = (organizationId, month) => {
43+
if (!mongoose.Types.ObjectId.isValid(organizationId)) return null;
44+
return BillingUsage.findOneAndUpdate(
45+
{ organizationId, month },
46+
{ $set: { counters: {} } },
47+
{ returnDocument: 'after' },
48+
).exec();
49+
};
50+
51+
export default {
52+
get,
53+
increment,
54+
reset,
55+
};
Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
/**
2+
* Module dependencies
3+
*/
4+
import UsageRepository from '../repositories/billing.usage.repository.js';
5+
6+
/**
7+
* Compute the current month string in YYYY-MM format.
8+
* @returns {String} e.g. '2026-03'
9+
*/
10+
const currentMonth = () => {
11+
const now = new Date();
12+
const year = now.getFullYear();
13+
const month = String(now.getMonth() + 1).padStart(2, '0');
14+
return `${year}-${month}`;
15+
};
16+
17+
/**
18+
* @desc Increment a usage counter for the given organization (current month).
19+
* @param {String} organizationId - The organization ID.
20+
* @param {String} key - The counter key to increment.
21+
* @param {Number} amount - The amount to increment by.
22+
* @returns {Promise<Object>} The updated usage document.
23+
*/
24+
const increment = (organizationId, key, amount) => UsageRepository.increment(organizationId, currentMonth(), key, amount);
25+
26+
/**
27+
* @desc Get usage for the given organization (current month).
28+
* @param {String} organizationId - The organization ID.
29+
* @returns {Promise<Object>} The usage document or an object with empty counters.
30+
*/
31+
const get = async (organizationId) => {
32+
const usage = await UsageRepository.get(organizationId, currentMonth());
33+
return usage || { organizationId, month: currentMonth(), counters: {} };
34+
};
35+
36+
/**
37+
* @desc Reset usage counters for the given organization (current month).
38+
* @param {String} organizationId - The organization ID.
39+
* @returns {Promise<Object|null>} The updated usage document or null.
40+
*/
41+
const reset = (organizationId) => UsageRepository.reset(organizationId, currentMonth());
42+
43+
export default {
44+
increment,
45+
get,
46+
reset,
47+
};
Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
/**
2+
* Module dependencies.
3+
*/
4+
import { jest, beforeEach, afterEach } from '@jest/globals';
5+
import schema from '../models/billing.usage.schema.js';
6+
7+
/**
8+
* Unit tests
9+
*/
10+
describe('BillingUsage unit tests:', () => {
11+
describe('Schema validation', () => {
12+
let usage;
13+
14+
beforeEach(() => {
15+
usage = {
16+
organizationId: '507f1f77bcf86cd799439011',
17+
month: '2026-03',
18+
counters: { executions: 10 },
19+
};
20+
});
21+
22+
test('should be valid with correct data', () => {
23+
const result = schema.BillingUsage.safeParse(usage);
24+
expect(result.error).toBeFalsy();
25+
expect(result.data.organizationId).toBe('507f1f77bcf86cd799439011');
26+
expect(result.data.month).toBe('2026-03');
27+
expect(result.data.counters.executions).toBe(10);
28+
});
29+
30+
test('should show error when organizationId is missing', () => {
31+
usage.organizationId = '';
32+
const result = schema.BillingUsage.safeParse(usage);
33+
expect(result.error).toBeDefined();
34+
});
35+
36+
test('should show error when organizationId is invalid', () => {
37+
usage.organizationId = 'not-valid';
38+
const result = schema.BillingUsage.safeParse(usage);
39+
expect(result.error).toBeDefined();
40+
});
41+
42+
test('should show error when month format is invalid', () => {
43+
usage.month = '2026/03';
44+
const result = schema.BillingUsage.safeParse(usage);
45+
expect(result.error).toBeDefined();
46+
});
47+
48+
test('should show error when month is missing', () => {
49+
delete usage.month;
50+
const result = schema.BillingUsage.safeParse(usage);
51+
expect(result.error).toBeDefined();
52+
});
53+
54+
test('should default counters to empty object', () => {
55+
delete usage.counters;
56+
const result = schema.BillingUsage.safeParse(usage);
57+
expect(result.error).toBeFalsy();
58+
expect(result.data.counters).toEqual({});
59+
});
60+
});
61+
62+
describe('Service layer', () => {
63+
let BillingUsageService;
64+
let mockUsageRepository;
65+
66+
const orgId = '507f1f77bcf86cd799439011';
67+
68+
beforeEach(async () => {
69+
jest.resetModules();
70+
71+
mockUsageRepository = {
72+
get: jest.fn(),
73+
increment: jest.fn(),
74+
reset: jest.fn(),
75+
};
76+
77+
jest.unstable_mockModule('../repositories/billing.usage.repository.js', () => ({
78+
default: mockUsageRepository,
79+
}));
80+
81+
const mod = await import('../services/billing.usage.service.js');
82+
BillingUsageService = mod.default;
83+
});
84+
85+
afterEach(() => {
86+
jest.restoreAllMocks();
87+
});
88+
89+
test('increment should create document on first call (upsert)', async () => {
90+
const created = { organizationId: orgId, month: '2026-03', counters: { executions: 1 } };
91+
mockUsageRepository.increment.mockResolvedValue(created);
92+
93+
const result = await BillingUsageService.increment(orgId, 'executions', 1);
94+
95+
expect(mockUsageRepository.increment).toHaveBeenCalledWith(orgId, expect.stringMatching(/^\d{4}-\d{2}$/), 'executions', 1);
96+
expect(result.counters.executions).toBe(1);
97+
});
98+
99+
test('increment should atomically increase counter', async () => {
100+
const updated = { organizationId: orgId, month: '2026-03', counters: { executions: 5 } };
101+
mockUsageRepository.increment.mockResolvedValue(updated);
102+
103+
const result = await BillingUsageService.increment(orgId, 'executions', 3);
104+
105+
expect(result.counters.executions).toBe(5);
106+
});
107+
108+
test('get should return empty counters for new org/month', async () => {
109+
mockUsageRepository.get.mockResolvedValue(null);
110+
111+
const result = await BillingUsageService.get(orgId);
112+
113+
expect(result.counters).toEqual({});
114+
expect(result.organizationId).toBe(orgId);
115+
});
116+
117+
test('get should return existing usage document', async () => {
118+
const existing = { organizationId: orgId, month: '2026-03', counters: { executions: 42, aiCalls: 5 } };
119+
mockUsageRepository.get.mockResolvedValue(existing);
120+
121+
const result = await BillingUsageService.get(orgId);
122+
123+
expect(result.counters.executions).toBe(42);
124+
expect(result.counters.aiCalls).toBe(5);
125+
});
126+
127+
test('reset should clear counters', async () => {
128+
const resetDoc = { organizationId: orgId, month: '2026-03', counters: {} };
129+
mockUsageRepository.reset.mockResolvedValue(resetDoc);
130+
131+
const result = await BillingUsageService.reset(orgId);
132+
133+
expect(mockUsageRepository.reset).toHaveBeenCalledWith(orgId, expect.stringMatching(/^\d{4}-\d{2}$/));
134+
expect(result.counters).toEqual({});
135+
});
136+
});
137+
});

0 commit comments

Comments
 (0)