Skip to content

Commit d36a92d

Browse files
brandonkachenbrandonkachen
committed
test: add unit tests for calculateTotalLegacyReferralBonus and one-time referral grants
1 parent ac94988 commit d36a92d

File tree

2 files changed

+684
-0
lines changed

2 files changed

+684
-0
lines changed

‎packages/billing/src/__tests__/grant-credits.test.ts‎

Lines changed: 313 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,180 @@ describe('grant-credits', () => {
8888
clearMockedModules()
8989
})
9090

91+
describe('calculateTotalLegacyReferralBonus', () => {
92+
const createDbMockForReferralQuery = (totalCredits: string | null) => ({
93+
select: () => ({
94+
from: () => ({
95+
where: () => Promise.resolve([{ totalCredits }]),
96+
}),
97+
}),
98+
})
99+
100+
const createDbMockThatThrows = (error: Error) => ({
101+
select: () => ({
102+
from: () => ({
103+
where: () => Promise.reject(error),
104+
}),
105+
}),
106+
})
107+
108+
it('should return total credits when user has legacy referrals as referrer', async () => {
109+
await mockModule('@codebuff/internal/db', () => ({
110+
default: createDbMockForReferralQuery('500'),
111+
}))
112+
113+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
114+
115+
const result = await calculateTotalLegacyReferralBonus({
116+
userId: 'user-123',
117+
logger,
118+
})
119+
120+
expect(result).toBe(500)
121+
})
122+
123+
it('should return total credits when user has legacy referrals as referred', async () => {
124+
await mockModule('@codebuff/internal/db', () => ({
125+
default: createDbMockForReferralQuery('250'),
126+
}))
127+
128+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
129+
130+
const result = await calculateTotalLegacyReferralBonus({
131+
userId: 'referred-user',
132+
logger,
133+
})
134+
135+
expect(result).toBe(250)
136+
})
137+
138+
it('should return combined total when user has legacy referrals as both referrer and referred', async () => {
139+
await mockModule('@codebuff/internal/db', () => ({
140+
default: createDbMockForReferralQuery('750'),
141+
}))
142+
143+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
144+
145+
const result = await calculateTotalLegacyReferralBonus({
146+
userId: 'user-with-both',
147+
logger,
148+
})
149+
150+
expect(result).toBe(750)
151+
})
152+
153+
it('should return 0 when user has no legacy referrals (only non-legacy)', async () => {
154+
// The query filters by is_legacy = true, so non-legacy referrals return 0
155+
await mockModule('@codebuff/internal/db', () => ({
156+
default: createDbMockForReferralQuery('0'),
157+
}))
158+
159+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
160+
161+
const result = await calculateTotalLegacyReferralBonus({
162+
userId: 'user-with-only-new-referrals',
163+
logger,
164+
})
165+
166+
expect(result).toBe(0)
167+
})
168+
169+
it('should return 0 when user has no referrals at all', async () => {
170+
await mockModule('@codebuff/internal/db', () => ({
171+
default: createDbMockForReferralQuery('0'),
172+
}))
173+
174+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
175+
176+
const result = await calculateTotalLegacyReferralBonus({
177+
userId: 'user-with-no-referrals',
178+
logger,
179+
})
180+
181+
expect(result).toBe(0)
182+
})
183+
184+
it('should return 0 when query returns null (COALESCE handles this)', async () => {
185+
await mockModule('@codebuff/internal/db', () => ({
186+
default: createDbMockForReferralQuery(null),
187+
}))
188+
189+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
190+
191+
const result = await calculateTotalLegacyReferralBonus({
192+
userId: 'user-null-result',
193+
logger,
194+
})
195+
196+
expect(result).toBe(0)
197+
})
198+
199+
it('should return 0 when query returns undefined result', async () => {
200+
await mockModule('@codebuff/internal/db', () => ({
201+
default: {
202+
select: () => ({
203+
from: () => ({
204+
where: () => Promise.resolve([]),
205+
}),
206+
}),
207+
},
208+
}))
209+
210+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
211+
212+
const result = await calculateTotalLegacyReferralBonus({
213+
userId: 'user-empty-result',
214+
logger,
215+
})
216+
217+
expect(result).toBe(0)
218+
})
219+
220+
it('should return 0 and log error when database query fails', async () => {
221+
const dbError = new Error('Database connection failed')
222+
await mockModule('@codebuff/internal/db', () => ({
223+
default: createDbMockThatThrows(dbError),
224+
}))
225+
226+
const errorLogs: any[] = []
227+
const errorLogger: Logger = {
228+
...logger,
229+
error: (...args: any[]) => {
230+
errorLogs.push(args)
231+
},
232+
}
233+
234+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
235+
236+
const result = await calculateTotalLegacyReferralBonus({
237+
userId: 'user-db-error',
238+
logger: errorLogger,
239+
})
240+
241+
expect(result).toBe(0)
242+
expect(errorLogs.length).toBe(1)
243+
expect(errorLogs[0][0]).toMatchObject({
244+
userId: 'user-db-error',
245+
error: dbError,
246+
})
247+
})
248+
249+
it('should handle large credit values correctly', async () => {
250+
await mockModule('@codebuff/internal/db', () => ({
251+
default: createDbMockForReferralQuery('999999'),
252+
}))
253+
254+
const { calculateTotalLegacyReferralBonus } = await import('../grant-credits')
255+
256+
const result = await calculateTotalLegacyReferralBonus({
257+
userId: 'power-referrer',
258+
logger,
259+
})
260+
261+
expect(result).toBe(999999)
262+
})
263+
})
264+
91265
describe('triggerMonthlyResetAndGrant', () => {
92266
describe('autoTopupEnabled return value', () => {
93267
it('should return autoTopupEnabled: true when user has auto_topup_enabled: true', async () => {
@@ -200,5 +374,144 @@ describe('grant-credits', () => {
200374
expect(result.quotaResetDate).toEqual(futureDate)
201375
})
202376
})
377+
378+
describe('legacy referral grants', () => {
379+
// Track grant operations to verify type and expiration
380+
let grantCalls: any[] = []
381+
382+
const createTxMockWithGrants = (user: {
383+
next_quota_reset: Date | null
384+
auto_topup_enabled: boolean | null
385+
} | null, legacyReferralBonus: number) => {
386+
grantCalls = []
387+
return {
388+
query: {
389+
user: {
390+
findFirst: async () => user,
391+
},
392+
},
393+
update: () => ({
394+
set: () => ({
395+
where: () => Promise.resolve(),
396+
}),
397+
}),
398+
insert: () => ({
399+
values: (values: any) => {
400+
grantCalls.push(values)
401+
return {
402+
onConflictDoNothing: () => ({
403+
returning: () => Promise.resolve([{ id: 'test-id' }]),
404+
}),
405+
}
406+
},
407+
}),
408+
select: () => ({
409+
from: () => ({
410+
where: () => ({
411+
orderBy: () => ({
412+
limit: () => [],
413+
}),
414+
}),
415+
then: (cb: any) => cb([{ totalCredits: String(legacyReferralBonus) }]),
416+
}),
417+
}),
418+
execute: () => Promise.resolve([]),
419+
}
420+
}
421+
422+
const createTransactionMockWithGrants = (user: {
423+
next_quota_reset: Date | null
424+
auto_topup_enabled: boolean | null
425+
} | null, legacyReferralBonus: number) => ({
426+
withAdvisoryLockTransaction: async ({
427+
callback,
428+
}: {
429+
callback: (tx: any) => Promise<any>
430+
}) => ({ result: await callback(createTxMockWithGrants(user, legacyReferralBonus)), lockWaitMs: 0 }),
431+
})
432+
433+
it('should grant referral_legacy type when user has legacy referrals and quota needs reset', async () => {
434+
const pastResetDate = new Date(Date.now() - 24 * 60 * 60 * 1000) // Yesterday
435+
const user = {
436+
next_quota_reset: pastResetDate,
437+
auto_topup_enabled: false,
438+
}
439+
const legacyReferralBonus = 500
440+
441+
await mockModule('@codebuff/internal/db', () => ({
442+
default: {
443+
select: () => ({
444+
from: () => ({
445+
where: () => ({
446+
orderBy: () => ({
447+
limit: () => [],
448+
}),
449+
}),
450+
}),
451+
}),
452+
},
453+
}))
454+
await mockModule('@codebuff/internal/db/transaction', () =>
455+
createTransactionMockWithGrants(user, legacyReferralBonus),
456+
)
457+
458+
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
459+
460+
await fn({
461+
userId: 'user-with-legacy-referrals',
462+
logger,
463+
})
464+
465+
// Should have made 2 grant calls (free + referral_legacy)
466+
expect(grantCalls.length).toBe(2)
467+
468+
// Find the referral grant
469+
const referralGrant = grantCalls.find((call) => call.type === 'referral_legacy')
470+
expect(referralGrant).toBeDefined()
471+
expect(referralGrant.principal).toBe(legacyReferralBonus)
472+
expect(referralGrant.balance).toBe(legacyReferralBonus)
473+
expect(referralGrant.expires_at).toBeDefined() // Legacy referrals expire at next reset
474+
expect(referralGrant.description).toBe('Monthly referral bonus (legacy)')
475+
})
476+
477+
it('should NOT grant referral credits when user has no legacy referrals', async () => {
478+
const pastResetDate = new Date(Date.now() - 24 * 60 * 60 * 1000) // Yesterday
479+
const user = {
480+
next_quota_reset: pastResetDate,
481+
auto_topup_enabled: false,
482+
}
483+
const legacyReferralBonus = 0 // No legacy referrals
484+
485+
await mockModule('@codebuff/internal/db', () => ({
486+
default: {
487+
select: () => ({
488+
from: () => ({
489+
where: () => ({
490+
orderBy: () => ({
491+
limit: () => [],
492+
}),
493+
}),
494+
}),
495+
}),
496+
},
497+
}))
498+
await mockModule('@codebuff/internal/db/transaction', () =>
499+
createTransactionMockWithGrants(user, legacyReferralBonus),
500+
)
501+
502+
const { triggerMonthlyResetAndGrant: fn } = await import('../grant-credits')
503+
504+
await fn({
505+
userId: 'user-without-legacy-referrals',
506+
logger,
507+
})
508+
509+
// Should only have made 1 grant call (free only, no referral)
510+
expect(grantCalls.length).toBe(1)
511+
512+
// The only grant should be 'free' type
513+
expect(grantCalls[0].type).toBe('free')
514+
})
515+
})
203516
})
204517
})

0 commit comments

Comments
 (0)