fix left overs

Author: npalm

lambdas/functions/control-plane/src/lambda.test.ts

@@ -1,6 +1,5 @@
 import { captureLambdaHandler, logger } from '@aws-github-runner/aws-powertools-util';
 import { Context, SQSEvent, SQSRecord } from 'aws-lambda';
-import { mocked } from 'jest-mock';
 
 import { addMiddleware, adjustPool, scaleDownHandler, scaleUpHandler, ssmHousekeeper, jobRetryCheck } from './lambda';
 import { adjust } from './pool/pool';
@@ -80,7 +79,7 @@ describe('Test scale up lambda wrapper.', () => {
   });
 
   it('Scale without error should resolve.', async () => {
-    const mock = mocked(scaleUp);
+    const mock = vi.fn(scaleUp);
     mock.mockImplementation(() => {
       return new Promise((resolve) => {
         resolve();
@@ -91,21 +90,24 @@ describe('Test scale up lambda wrapper.', () => {
 
   it('Non scale should resolve.', async () => {
     const error = new Error('Non scale should resolve.');
-    const mock = mocked(scaleUp);
+    const mock = vi.fn(scaleUp);
     mock.mockRejectedValue(error);
-    await expect(scaleUpHandler(sqsEvent, context)).resolves.not.toThrow;
+    await expect(scaleUpHandler(sqsEvent, context)).resolves.not.toThrow();
   });
 
   it('Scale should be rejected', async () => {
     const error = new ScaleError('Scale should be rejected');
-    const mock = mocked(scaleUp);
-    mock.mockRejectedValue(error);
+    const mock = vi.fn() as MockedFunction<typeof scaleUp>;
+    mock.mockImplementation(() => {
+      return Promise.reject(error);
+    });
+    vi.mocked(scaleUp).mockImplementation(mock);
     await expect(scaleUpHandler(sqsEvent, context)).rejects.toThrow(error);
   });
 });
 
 async function testInvalidRecords(sqsRecords: SQSRecord[]) {
-  const mock = mocked(scaleUp);
+  const mock = vi.fn(scaleUp);
   const logWarnSpy = vi.spyOn(logger, 'warn');
   mock.mockImplementation(() => {
     return new Promise((resolve) => {
@@ -127,7 +129,7 @@ async function testInvalidRecords(sqsRecords: SQSRecord[]) {
 
 describe('Test scale down lambda wrapper.', () => {
   it('Scaling down no error.', async () => {
-    const mock = mocked(scaleDown);
+    const mock = vi.fn(scaleDown);
     mock.mockImplementation(() => {
       return new Promise((resolve) => {
         resolve();
@@ -138,15 +140,15 @@ describe('Test scale down lambda wrapper.', () => {
 
   it('Scaling down with error.', async () => {
     const error = new Error('Scaling down with error.');
-    const mock = mocked(scaleDown);
+    const mock = vi.fn(scaleDown);
     mock.mockRejectedValue(error);
     await expect(scaleDownHandler({}, context)).resolves.not.toThrow();
   });
 });
 
 describe('Adjust pool.', () => {
   it('Receive message to adjust pool.', async () => {
-    const mock = mocked(adjust);
+    const mock = vi.fn(adjust);
     mock.mockImplementation(() => {
       return new Promise((resolve) => {
         resolve();
@@ -156,12 +158,15 @@ describe('Adjust pool.', () => {
   });
 
   it('Handle error for adjusting pool.', async () => {
-    const mock = mocked(adjust);
     const error = new Error('Handle error for adjusting pool.');
-    mock.mockRejectedValue(error);
+    const mock = vi.fn() as MockedFunction<typeof adjust>;
+    mock.mockImplementation(() => {
+      return Promise.reject(error);
+    });
+    vi.mocked(adjust).mockImplementation(mock);
     const logSpy = vi.spyOn(logger, 'error');
     await adjustPool({ poolSize: 0 }, context);
-    expect(logSpy).lastCalledWith(expect.stringContaining(error.message), expect.anything());
+    expect(logSpy).toHaveBeenCalledWith(`Handle error for adjusting pool. ${error.message}`, { error });
   });
 });
 
@@ -175,7 +180,7 @@ describe('Test middleware', () => {
 
 describe('Test ssm housekeeper lambda wrapper.', () => {
   it('Invoke without errors.', async () => {
-    const mock = mocked(cleanSSMTokens);
+    const mock = vi.fn(cleanSSMTokens);
     mock.mockImplementation(() => {
       return new Promise((resolve) => {
         resolve();
@@ -192,31 +197,32 @@ describe('Test ssm housekeeper lambda wrapper.', () => {
   });
 
   it('Errors not throws.', async () => {
-    const mock = mocked(cleanSSMTokens);
+    const mock = vi.fn(cleanSSMTokens);
     mock.mockRejectedValue(new Error());
     await expect(ssmHousekeeper({}, context)).resolves.not.toThrow();
   });
 });
 
 describe('Test job retry check wrapper', () => {
   it('Handle without error should resolve.', async () => {
-    const mock = mocked(checkAndRetryJob);
+    const mock = vi.fn() as MockedFunction<typeof checkAndRetryJob>;
     mock.mockImplementation(() => {
-      return new Promise((resolve) => {
-        resolve();
-      });
+      return Promise.resolve();
     });
+    vi.mocked(checkAndRetryJob).mockImplementation(mock);
     await expect(jobRetryCheck(sqsEvent, context)).resolves.not.toThrow();
   });
 
   it('Handle with error should resolve and log only a warning.', async () => {
-    const logSpyWarn = vi.spyOn(logger, 'warn');
-
-    const mock = mocked(checkAndRetryJob);
     const error = new Error('Error handling retry check.');
-    mock.mockRejectedValue(error);
+    const mock = vi.fn() as MockedFunction<typeof checkAndRetryJob>;
+    mock.mockImplementation(() => {
+      return Promise.reject(error);
+    });
+    vi.mocked(checkAndRetryJob).mockImplementation(mock);
 
+    const logSpyWarn = vi.spyOn(logger, 'warn');
     await expect(jobRetryCheck(sqsEvent, context)).resolves.not.toThrow();
-    expect(logSpyWarn).toHaveBeenCalledWith(expect.stringContaining(error.message), expect.anything());
+    expect(logSpyWarn).toHaveBeenCalledWith(`Error processing job retry: ${error.message}`, { error });
   });
 });

lambdas/functions/control-plane/src/lambda.ts

@@ -16,16 +16,18 @@ export async function scaleUpHandler(event: SQSEvent, context: Context): Promise
 
   if (event.Records.length !== 1) {
     logger.warn('Event ignored, only one record at the time can be handled, ensure the lambda batch size is set to 1.');
-    return new Promise((resolve) => resolve());
+    return Promise.resolve();
   }
 
   try {
     await scaleUp(event.Records[0].eventSource, JSON.parse(event.Records[0].body));
+    return Promise.resolve();
   } catch (e) {
     if (e instanceof ScaleError) {
-      throw e;
+      return Promise.reject(e);
     } else {
       logger.warn(`Ignoring error: ${e}`);
+      return Promise.resolve();
     }
   }
 }
@@ -48,8 +50,9 @@ export async function adjustPool(event: PoolEvent, context: Context): Promise<vo
   try {
     await adjust(event);
   } catch (e) {
-    logger.error(`${(e as Error).message}`, { error: e as Error });
+    logger.error(`Handle error for adjusting pool. ${(e as Error).message}`, { error: e as Error });
   }
+  return Promise.resolve();
 }
 
 export const addMiddleware = () => {
@@ -86,4 +89,5 @@ export async function jobRetryCheck(event: SQSEvent, context: Context): Promise<
       logger.warn(`Error processing job retry: ${e.message}`, { error: e });
     });
   }
+  return Promise.resolve();
 }