test: add more tests

Author: alexey-yarmosh

Diff for: src/lib/credits.ts

@@ -7,7 +7,7 @@ const ER_CONSTRAINT_FAILED_CODE = 4025;
 export class Credits {
 	constructor (private readonly sql: Knex) {}
 
-	async consume (userId: string, credits: number): Promise<{ isConsumed: boolean, remainingCredits: number }> {
+	async consume (userId: string, credits: number): Promise<{ isConsumed: boolean, remainingCredits?: number }> {
 		try {
 			const result = await this.sql.raw<[[{amount: number | null}]]>(`
 				INSERT INTO ?? (user_id, amount)
@@ -20,27 +20,18 @@ export class Credits {
 			const remainingCredits = result[0]?.[0]?.amount;
 
 			if (remainingCredits || remainingCredits === 0) {
-				return this.returnSuccess(remainingCredits);
+				return { isConsumed: true, remainingCredits };
 			}
 
-			return this.returnFail(userId);
+			return { isConsumed: false };
 		} catch (error) {
 			if (error && (error as Error & {errno?: number}).errno === ER_CONSTRAINT_FAILED_CODE) {
-				return this.returnFail(userId);
+				return { isConsumed: false };
 			}
 
 			throw error;
 		}
 	}
-
-	returnSuccess (remainingCredits: number) {
-		return { isConsumed: true, remainingCredits };
-	}
-
-	async returnFail (userId: string) {
-		const result = await this.sql(CREDITS_TABLE).where({ user_id: userId }).select<[{amount: number}]>('amount');
-		return { isConsumed: false, remainingCredits: result[0]?.amount || 0 };
-	}
 }
 
 export const credits = new Credits(client);

Diff for: src/lib/rate-limiter.ts

@@ -44,16 +44,14 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
 	} catch (error) {
 		if (error instanceof RateLimiterRes) {
 			if (ctx.state.userId) {
-				const { isConsumed, consumedCredits, remainingCredits, pointsToReward } = await consumeCredits(ctx, error, numberOfProbes);
+				const { isConsumed, consumedCredits, remainingCredits, pointsToReward } = await consumeCredits(ctx.state.userId, error, numberOfProbes);
 
 				if (isConsumed) {
 					const result = await rateLimiter.reward(id, pointsToReward);
-					setCreditsHeaders(ctx, consumedCredits, remainingCredits);
+					setCreditsHeaders(ctx, consumedCredits!, remainingCredits!);
 					setRateLimitHeaders(ctx, result, rateLimiter);
 					return;
 				}
-
-				setCreditsHeaders(ctx, consumedCredits, remainingCredits);
 			}
 
 			const result = await rateLimiter.reward(id, numberOfProbes);
@@ -65,12 +63,12 @@ export const rateLimit = async (ctx: ExtendedContext, numberOfProbes: number) =>
 	}
 };
 
-const consumeCredits = async (ctx: ExtendedContext, rateLimiterRes: RateLimiterRes, numberOfProbes: number) => {
+const consumeCredits = async (userId: string, rateLimiterRes: RateLimiterRes, numberOfProbes: number) => {
 	const freePoints = config.get<number>('measurement.authenticatedRateLimit');
 	const alreadyUsedPoints = rateLimiterRes.consumedPoints - numberOfProbes;
 	const remainingFreePoints = freePoints - alreadyUsedPoints;
 	const requiredPoints = numberOfProbes - remainingFreePoints;
-	const { isConsumed, remainingCredits } = await credits.consume(ctx.state.userId!, requiredPoints);
+	const { isConsumed, remainingCredits } = await credits.consume(userId, requiredPoints);
 
 	if (isConsumed) {
 		return {
@@ -83,9 +81,6 @@ const consumeCredits = async (ctx: ExtendedContext, rateLimiterRes: RateLimiterR
 
 	return {
 		isConsumed: false,
-		consumedCredits: 0,
-		remainingCredits,
-		pointsToReward: 0,
 	};
 };
 

Diff for: test/tests/integration/ratelimit.test.ts

@@ -145,7 +145,7 @@ describe('rate limiter', () => {
 				target: 'jsdelivr.com',
 			}).expect(202) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(99999);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('99999');
 		});
 
 		it('should fail (limit reached)', async () => {
@@ -156,7 +156,7 @@ describe('rate limiter', () => {
 				target: 'jsdelivr.com',
 			}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(0);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
 		});
 
 		it('should consume all points successfully or none at all (cost > remaining > 0)', async () => {
@@ -168,7 +168,7 @@ describe('rate limiter', () => {
 				limit: 2,
 			}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(1);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('1');
 		});
 	});
 
@@ -183,7 +183,7 @@ describe('rate limiter', () => {
 					target: 'jsdelivr.com',
 				}).expect(202) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(249);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('249');
 		});
 
 		it('should fail (limit reached)', async () => {
@@ -196,7 +196,7 @@ describe('rate limiter', () => {
 					target: 'jsdelivr.com',
 				}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(0);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
 		});
 
 		it('should consume all points successfully or none at all (cost > remaining > 0)', async () => {
@@ -210,7 +210,7 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(1);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('1');
 		});
 	});
 
@@ -235,7 +235,9 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(202) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(248);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('248');
+			expect(response.headers['x-credits-cost']).to.not.exist;
+			expect(response.headers['x-credits-remaining']).to.not.exist;
 			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
 			expect(amount).to.equal(10);
 		});
@@ -251,7 +253,9 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(202) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(0);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
+			expect(response.headers['x-credits-cost']).to.equal('2');
+			expect((response.headers['x-credits-remaining'])).to.equal('8');
 			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
 			expect(amount).to.equal(8);
 		});
@@ -267,7 +271,9 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(202) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(0);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
+			expect(response.headers['x-credits-cost']).to.equal('1');
+			expect((response.headers['x-credits-remaining'])).to.equal('9');
 			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
 			expect(amount).to.equal(9);
 		});
@@ -289,7 +295,9 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(0);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
+			expect(response.headers['x-credits-cost']).to.not.exist;
+			expect(response.headers['x-credits-remaining']).to.not.exist;
 			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
 			expect(amount).to.equal(1);
 		});
@@ -311,9 +319,33 @@ describe('rate limiter', () => {
 					limit: 2,
 				}).expect(429) as Response;
 
-			expect(Number(response.headers['x-ratelimit-remaining'])).to.equal(1);
+			expect(response.headers['x-ratelimit-remaining']).to.equal('1');
+			expect(response.headers['x-credits-cost']).to.not.exist;
+			expect(response.headers['x-credits-remaining']).to.not.exist;
 			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
 			expect(amount).to.equal(0);
 		});
+
+		it('should work fine if there is no credits row for that user', async () => {
+			await authenticatedRateLimiter.set('89da69bd-a236-4ab7-9c5d-b5f52ce09959', 250, 0);
+
+			await client(CREDITS_TABLE).where({
+				user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959',
+			}).delete();
+
+			const response = await requestAgent.post('/v1/measurements')
+				.set('Authorization', 'Bearer v2lUHEVLtVSskaRKDBabpyp4AkzdMnob')
+				.send({
+					type: 'ping',
+					target: 'jsdelivr.com',
+					limit: 2,
+				}).expect(429) as Response;
+
+			expect(response.headers['x-ratelimit-remaining']).to.equal('0');
+			expect(response.headers['x-credits-cost']).to.not.exist;
+			expect(response.headers['x-credits-remaining']).to.not.exist;
+			const [{ amount }] = await client(CREDITS_TABLE).select('amount').where({ user_id: '89da69bd-a236-4ab7-9c5d-b5f52ce09959' });
+			expect(amount).to.equal(null);
+		});
 	});
 });

Diff for: test/tests/unit/credits.test.ts

@@ -26,20 +26,18 @@ describe('Credits', () => {
 
 	it(`should return false if row wasn't updated`, async () => {
 		sqlStub.raw.resolves([ [{ amount: null }] ]);
-		selectStub.resolves([{ amount: null }]);
 		const credits = new Credits(sqlStub as unknown as Knex);
 		const result = await credits.consume('userId', 10);
-		expect(result).to.deep.equal({ isConsumed: false, remainingCredits: 0 });
+		expect(result).to.deep.equal({ isConsumed: false });
 	});
 
 	it(`should return false if update throws ER_CONSTRAINT_FAILED_CODE`, async () => {
 		const error: Error & {errno?: number} = new Error('constraint');
-		selectStub.resolves([{ amount: 5 }]);
 		error.errno = 4025;
 		sqlStub.raw.rejects(error);
 		const credits = new Credits(sqlStub as unknown as Knex);
 		const result = await credits.consume('userId', 10);
-		expect(result).to.deep.equal({ isConsumed: false, remainingCredits: 5 });
+		expect(result).to.deep.equal({ isConsumed: false });
 	});
 
 	it(`should throw if update throws other error`, async () => {