diff --git a/packages/account-usage/lib/usage.integration.test.ts b/packages/account-usage/lib/usage.integration.test.ts index 09f4d18fda7..23bfae0d85e 100644 --- a/packages/account-usage/lib/usage.integration.test.ts +++ b/packages/account-usage/lib/usage.integration.test.ts @@ -35,6 +35,9 @@ describe('Usage', () => { describe('get', () => { it('should return 0 for a non-existent-yet metric', async () => { + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + const res = (await usageTracker.get({ accountId: 1, metric: 'connections' })).unwrap(); expect(res).toEqual({ accountId: 1, metric: 'connections', current: 0 }); }); @@ -44,6 +47,9 @@ describe('Usage', () => { // Manually set an invalid entry in Redis await redis.set(`usageV2:${accountId}:${metric}`, 'invalid-json'); + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + const res = await usageTracker.get({ accountId, metric }); expect(res.isErr()).toBe(true); if (res.isErr()) { @@ -53,9 +59,56 @@ describe('Usage', () => { it('should return the current value for an existing metric', async () => { const accountId = 1; const metric = 'connections'; + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + await usageTracker.incr({ accountId, metric, delta: 5 }); + const res = (await usageTracker.get({ accountId, metric })).unwrap(); + expect(res).toEqual({ accountId, metric, current: 5 }); + }); + it('should trigger revalidation when entry is null', async () => { + const accountId = 1; + const metric = 'connections'; + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + const res = (await usageTracker.get({ accountId, metric })).unwrap(); + expect(res).toEqual({ accountId, metric, current: 0 }); + expect(revalidateSpy).toHaveBeenCalledTimes(1); + expect(revalidateSpy).toHaveBeenCalledWith({ accountId, metric }); + }); + it('should trigger revalidation when revalidateAfter has passed', async () => { + const accountId = 1; + const metric = 'connections'; + + // Set up an entry with a past revalidateAfter + await usageTracker.incr({ accountId, metric, delta: 5 }); + // Move time forward by 1 day to pass revalidateAfter + vi.advanceTimersByTime(24 * 60 * 60 * 1000); + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + const res = (await usageTracker.get({ accountId, metric })).unwrap(); + expect(res).toEqual({ accountId, metric, current: 5 }); + expect(revalidateSpy).toHaveBeenCalledTimes(1); + expect(revalidateSpy).toHaveBeenCalledWith({ accountId, metric }); + }); + it('should not trigger revalidation when entry exists and is not stale', async () => { + const accountId = 1; + const metric = 'connections'; + await usageTracker.incr({ accountId, metric, delta: 5 }); + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + const res = (await usageTracker.get({ accountId, metric })).unwrap(); expect(res).toEqual({ accountId, metric, current: 5 }); + // revalidateAfter hasn't passed yet + expect(revalidateSpy).not.toHaveBeenCalled(); }); }); @@ -85,8 +138,61 @@ describe('Usage', () => { it('should increment monthly metric', async () => { const accountId = 2; const metric = 'proxy'; + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + const res = (await usageTracker.incr({ accountId, metric, delta: 1 })).unwrap(); expect(res).toEqual({ accountId, metric, current: 1 }); }); }); + + describe('getAll', () => { + it('should trigger revalidation for null entries', async () => { + const accountId = 1; + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + const res = (await usageTracker.getAll(accountId)).unwrap(); + expect(res).toBeDefined(); + // Should trigger revalidation for all metrics that are null + expect(revalidateSpy).toHaveBeenCalled(); + }); + it('should trigger revalidation for stale entries', async () => { + const accountId = 1; + const metric = 'connections'; + + // Set up an entry with a past revalidateAfter + await usageTracker.incr({ accountId, metric, delta: 5 }); + // Move time forward by 1 day to pass revalidateAfter + vi.advanceTimersByTime(24 * 60 * 60 * 1000); + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + const res = (await usageTracker.getAll(accountId)).unwrap(); + expect(res).toBeDefined(); + expect(res[metric]).toEqual({ accountId, metric, current: 5 }); + // Should trigger revalidation for the stale metric + expect(revalidateSpy).toHaveBeenCalledWith({ accountId, metric }); + }); + it('should not trigger revalidation for non-stale entries', async () => { + const accountId = 1; + const metric = 'connections'; + + await usageTracker.incr({ accountId, metric, delta: 5 }); + + const revalidateSpy = vi.spyOn(usageTracker, 'revalidate'); + revalidateSpy.mockReturnValue(Promise.resolve(Ok(undefined))); + + const res = (await usageTracker.getAll(accountId)).unwrap(); + expect(res).toBeDefined(); + expect(res[metric]).toEqual({ accountId, metric, current: 5 }); + // revalidateAfter hasn't passed yet, so revalidation should not be called for this metric + // (but may be called for other null metrics) + const callsForMetric = revalidateSpy.mock.calls.filter((call) => call[0].metric === metric); + expect(callsForMetric).toHaveLength(0); + }); + }); }); diff --git a/packages/account-usage/lib/usage.ts b/packages/account-usage/lib/usage.ts index 3973c6fff5c..3d87a151561 100644 --- a/packages/account-usage/lib/usage.ts +++ b/packages/account-usage/lib/usage.ts @@ -88,6 +88,9 @@ export class UsageTracker implements IUsageTracker { if (entry.isErr()) { return Err(entry.error); } + if (entry.value === null || entry.value.revalidateAfter < now.getTime()) { + void this.revalidate({ accountId, metric }); + } return Ok({ accountId, metric, @@ -105,6 +108,9 @@ export class UsageTracker implements IUsageTracker { if (entry.isErr()) { return; } + if (entry.value === null || entry.value.revalidateAfter < now.getTime()) { + void this.revalidate({ accountId, metric: metric as UsageMetric }); + } result[metric as UsageMetric] = { accountId, metric: metric as UsageMetric,