@squiz/job-runner-lib 2.3.0 → 2.4.0

package/src/domains/job/model/JobExecution/mocks/JobExecutionDto.mock.ts

@@ -3,46 +3,57 @@
  * Copyright Squiz Australia Pty Ltd. All Rights Reserved.
  */
 
-import { JOB_EXECUTION_STATUS } from '../../../constants/JobExecutionStatus.constants';
-import { PARTITION_KEYS } from '../../../constants/Repository.constants';
 import { faker } from '@faker-js/faker/locale/en';
+import { JOB_EXECUTION_STATUS } from '../../../constants/JobExecutionStatus.constants';
+import { ENTITY_TYPES } from '../../../constants/Repository.constants';
 import { JobExecutionDto, JobExecutionRecord } from '../JobExecution';
 
 export const createMockJobExecutionDto = (partialRecord: Partial<JobExecutionRecord> = {}): JobExecutionDto => {
-  const jobName = faker.lorem.word();
+  const tenant = faker.string.uuid();
+  const jobName = partialRecord.jobName ?? faker.lorem.word();
   const timeQueued = faker.date.recent().toISOString();
   const id = `${timeQueued}-${faker.string.uuid()}`;
   const version = Array.from({ length: 3 })
     .map(() => faker.number.int())
     .join('.');
-  const input = faker.helpers.arrayElement([
-    faker.datatype.boolean(),
-    faker.lorem.word(),
-    faker.number.int(),
-    { [faker.lorem.word()]: faker.lorem.words },
-    null,
-    undefined,
-  ]);
+  const status = partialRecord.status ?? faker.helpers.arrayElement(Object.values(JOB_EXECUTION_STATUS));
+  const input =
+    'input' in partialRecord
+      ? partialRecord.input
+      : faker.helpers.arrayElement([
+          faker.datatype.boolean(),
+          faker.lorem.word(),
+          faker.number.int(),
+          { [faker.lorem.word()]: faker.lorem.words },
+          null,
+          undefined,
+        ]);
 
   return {
     context: {
       S: partialRecord.context ?? faker.string.uuid(),
     },
     ecsTaskId: partialRecord.ecsTaskId ? { S: partialRecord.ecsTaskId } : { NULL: true },
+    gsi1pk: {
+      S: `${ENTITY_TYPES.jobExecution}~${status}`,
+    },
     id: {
       S: partialRecord.id ?? id,
     },
     jobName: {
-      S: partialRecord.jobName ?? jobName,
+      S: jobName,
+    },
+    lsi1sk: {
+      S: `${status}~${jobName}`,
     },
     pk: {
-      S: partialRecord.pk ?? `${PARTITION_KEYS.jobExecution}-${partialRecord.jobName ?? jobName}`,
+      S: partialRecord.pk ?? `${ENTITY_TYPES.jobExecution}~${tenant}`,
     },
     sk: {
-      S: partialRecord.sk ?? `${timeQueued}-${id}`,
+      S: partialRecord.sk ?? `${jobName}~${id}`,
     },
     status: {
-      S: partialRecord.status ?? faker.helpers.arrayElement(Object.values(JOB_EXECUTION_STATUS)),
+      S: status,
     },
     timeFinished: partialRecord.timeFinished ? { S: partialRecord.timeFinished } : { NULL: true },
     timeQueued: {
@@ -52,7 +63,12 @@ export const createMockJobExecutionDto = (partialRecord: Partial<JobExecutionRec
     version: {
       S: partialRecord.version ?? version,
     },
-    input: input ? { S: JSON.stringify(input) } : { NULL: true },
+    input: input === undefined || input === null ? { NULL: true } : { S: JSON.stringify(input) },
+    retryCount: {
+      N: String(partialRecord.retryCount ?? 0),
+    },
+    lastFailureReason: partialRecord.lastFailureReason ? { S: partialRecord.lastFailureReason } : { NULL: true },
+    retryAfter: partialRecord.retryAfter ? { S: partialRecord.retryAfter } : { NULL: true },
   };
 };
 

package/src/domains/job/model/JobExecution/mocks/JobExecutionRecord.mock.ts

@@ -3,9 +3,9 @@
  * Copyright Squiz Australia Pty Ltd. All Rights Reserved.
  */
 
-import { JOB_EXECUTION_STATUS } from '../../../constants/JobExecutionStatus.constants';
-import { PARTITION_KEYS } from '../../../constants/Repository.constants';
 import { faker } from '@faker-js/faker/locale/en';
+import { JOB_EXECUTION_STATUS } from '../../../constants/JobExecutionStatus.constants';
+import { ENTITY_TYPES } from '../../../constants/Repository.constants';
 import { JobExecutionRecord } from '../JobExecution';
 
 export const createMockJobExecutionRecord = (partialRecord: Partial<JobExecutionRecord> = {}): JobExecutionRecord => {
@@ -15,6 +15,7 @@ export const createMockJobExecutionRecord = (partialRecord: Partial<JobExecution
   const version = Array.from({ length: 3 })
     .map(() => faker.number.int())
     .join('.');
+  const status = faker.helpers.objectValue(JOB_EXECUTION_STATUS);
   const input = faker.helpers.arrayElement([
     faker.datatype.boolean(),
     faker.lorem.word(),
@@ -27,16 +28,21 @@ export const createMockJobExecutionRecord = (partialRecord: Partial<JobExecution
   return {
     context: faker.string.uuid(),
     ecsTaskId: null,
+    gsi1pk: `${ENTITY_TYPES.jobExecution}~${status}`,
     id: id,
     jobName: jobName,
-    pk: `${PARTITION_KEYS.jobExecution}-${jobName}`,
+    lsi1sk: `${status}~${jobName}`,
+    pk: `${ENTITY_TYPES.jobExecution}-${jobName}`,
     sk: `${timeQueued}-${id}`,
-    status: faker.helpers.objectValue(JOB_EXECUTION_STATUS),
+    status: status,
     timeFinished: null,
     timeQueued: timeQueued,
     timeStarted: null,
     version,
     input,
+    retryCount: 0,
+    lastFailureReason: null,
+    retryAfter: null,
     ...partialRecord,
   };
 };

package/src/domains/job/model/LogPayload/LogPayload.ts

@@ -1,12 +1,33 @@
-export type LogPayload = {
-  message: string;
-  type?: string;
-  area?: string;
-  req_id?: string;
-  'log-command'?: string;
-  level?: string;
-  'function-name'?: string;
-  'job-name'?: string;
-  'job-version'?: string;
-  timestamp: string;
+import { z } from 'zod';
+
+export enum LogLevel {
+  AUDIT = 'AUDIT',
+  INFO = 'INFO',
+  WARNING = 'WARNING',
+  ERROR = 'ERROR',
+}
+
+export const LogPayloadSchema = z.object({
+  message: z.string(),
+  meta: z.object({
+    // Properties are nested under "meta" so the DXP log level does not conflict with the winston log level.
+    // Job runner specific attributes.
+    tenant: z.string(),
+    jobName: z.string(),
+    // Logging service attributes.
+    level: z.enum(LogLevel),
+    timestamp: z.iso.datetime(),
+    traceid: z.string(),
+    tags: z.array(z.string()).optional(),
+    details: z.record(z.string(), z.unknown()).optional(),
+  }),
+});
+
+export type LogPayload = z.infer<typeof LogPayloadSchema>;
+
+export const CONSOLE_LOG_METHOD_LEVELS: Partial<Record<keyof typeof console, LogLevel>> = {
+  log: LogLevel.INFO,
+  info: LogLevel.INFO,
+  warn: LogLevel.WARNING,
+  error: LogLevel.ERROR,
 };

package/src/domains/job/model/LogPayload/index.ts

@@ -1 +1,2 @@
 export * from './LogPayload';
+export * from './mocks/LogPayload.mock';

package/src/domains/job/model/LogPayload/mocks/LogPayload.mock.ts

@@ -0,0 +1,21 @@
+import { faker } from '@faker-js/faker';
+import { LogLevel, LogPayload } from '../LogPayload';
+
+type DeepPartial<T> = {
+  [K in keyof T]?: T[K] extends object ? DeepPartial<T[K]> : T[K];
+};
+
+export const createMockLogPayload = (partialRecord: DeepPartial<LogPayload>): LogPayload => {
+  return {
+    message: faker.lorem.words(),
+    ...partialRecord,
+    meta: {
+      traceid: faker.string.uuid(),
+      level: faker.helpers.objectValue(LogLevel),
+      timestamp: faker.date.recent().toISOString(),
+      tenant: faker.string.uuid(),
+      jobName: faker.string.uuid(),
+      ...partialRecord.meta,
+    },
+  };
+};

package/src/domains/job/repository/AbstractRepository.ts

@@ -6,21 +6,19 @@ import { injectable } from 'inversify';
 import 'reflect-metadata';
 import {
   AttributeValue,
+  BatchWriteItemCommand,
   DeleteItemCommand,
   DeleteItemCommandInput,
   DynamoDBClient,
   GetItemCommand,
   GetItemCommandInput,
   PutItemCommand,
-  ScanCommand,
-  ScanCommandInput,
-  ScanCommandOutput,
+  QueryCommand,
+  QueryCommandOutput,
   UpdateItemCommand,
   UpdateItemCommandInput,
 } from '@aws-sdk/client-dynamodb';
-import { parsePartitionKey } from '../ParsePartitionKey';
-import { PartitionKey } from '../constants/Repository.constants';
-import { BaseDto, BaseItem } from '../../job/model/Base/Base';
+import { BaseDto, BaseItem, ListParameters } from '../../job/model/Base/Base';
 
 /**
  * Generic abstract class for representing repositories interacting with DynamoDB.
@@ -31,22 +29,12 @@ import { BaseDto, BaseItem } from '../../job/model/Base/Base';
  * @template U The corresponding DB record attributes for the entity.
  */
 @injectable()
-/* eslint-disable-next-line @typescript-eslint/no-explicit-any */
-export abstract class Repository<T extends BaseItem, U extends BaseDto> {
-  protected pk: PartitionKey;
-  protected tableName: string;
-
-  public constructor(protected dynamoDbClient: DynamoDBClient) {
-    // HACK: workaround to address inheritance issue in InversifyJS
-    // Reference: https://github.com/inversify/InversifyJS/blob/master/wiki/inheritance.md
-
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore @ts-2322
-    this.pk = null;
-    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-    // @ts-ignore @ts-2322
-    this.tableName = null;
-  }
+export abstract class Repository<T extends BaseItem, U extends BaseDto, SK extends Partial<T>> {
+  public constructor(
+    protected dynamoDbClient: DynamoDBClient,
+    protected tableName: string,
+    protected pk: string,
+  ) {}
 
   /**
    * Save item into database.
@@ -55,7 +43,6 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
    * @returns {Promise<T>} The item stored.
    * @memberof Repository
    */
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
   async create(input: T): Promise<T> {
     const dynamoDbRecord = this.convertToDynamoDbItem(input);
 
@@ -76,11 +63,11 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
    * @returns {Promise<T | null>} The retrieved record.
    * @memberof Repository
    */
-  async findOne(sk: string): Promise<T | null> {
+  async findOne(sk: SK): Promise<T | null> {
     const params: GetItemCommandInput = {
       Key: {
-        pk: { S: parsePartitionKey(this.pk) },
-        sk: { S: sk },
+        pk: { S: this.pk },
+        sk: { S: this.buildSk(sk) },
       },
       TableName: this.tableName,
     };
@@ -94,41 +81,41 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
    * @returns {Promise<Array<T>>} The retrieved records.
    * @memberof Repository
    */
-  async list(): Promise<Array<T>> {
-    const params: ScanCommandInput = {
-      ExpressionAttributeNames: {
-        '#pk': 'pk',
-      },
-      ExpressionAttributeValues: {
-        ':pk': { S: this.pk },
-      },
-      FilterExpression: '#pk = :pk',
-      TableName: this.tableName,
-    };
-
+  async list(parameters: ListParameters<never> = {}): Promise<Array<T>> {
     // Apply pagination for the scan.
     // By default, the Scan operation processes data sequentially.
     // DynamoDB returns data in increments of 1 MB increments.
 
-    let startKey; /* Indicator of where to paginate from. */
+    let startKey: Record<string, AttributeValue> | undefined = undefined; /* Indicator of where to paginate from. */
     const results: Array<Record<string, AttributeValue>> = [];
 
     // Iterate through pages until no more results are found.
-    while (startKey !== null) {
-      const res: ScanCommandOutput = await this.dynamoDbClient.send(
-        new ScanCommand({ ...params, ExclusiveStartKey: startKey }),
+    do {
+      const res: QueryCommandOutput = await this.dynamoDbClient.send(
+        new QueryCommand({
+          ExpressionAttributeNames: {
+            '#pk': 'pk',
+          },
+          ExpressionAttributeValues: {
+            ':pk': { S: this.pk },
+          },
+          KeyConditionExpression: '#pk = :pk',
+          TableName: this.tableName,
+          ExclusiveStartKey: startKey,
+          Limit: parameters.limit,
+          ConsistentRead: parameters.consistentRead,
+        }),
       );
 
-      if (!res) {
-        // Stop paginating if result is unexpected.
-        startKey = null;
-      } else {
-        results.push(...(res.Items ?? []));
+      results.push(...(res.Items ?? []));
+      startKey = res.LastEvaluatedKey;
 
-        // If no more results, set key to null to terminate loop
-        startKey = res.LastEvaluatedKey ?? null;
+      if (parameters.limit !== undefined && results.length >= parameters.limit) {
+        // Truncate list to max requested items, stop querying.
+        results.length = parameters.limit;
+        break;
       }
-    }
+    } while (startKey);
 
     return results.length ? results.map((item) => this.convertToItem(item as U)) : [];
   }
@@ -136,16 +123,16 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
   /**
    * Update record in database.
    * @param {string} sk The sort key used to identify the record.
-   * @param {Partial<T>} input The (partial) record to update.
+   * @param {Partial<T>} input The (partial) record to update (properties used in the SK cannot be updated).
    * @returns {Promise<T | null>} The updated record.
    * @memberof Repository
    */
-  async update(sk: string, input: Partial<T>): Promise<T | null> {
+  async update(sk: SK, input: Partial<Omit<T, keyof SK>>): Promise<T> {
     // Initialise update command
     const updateParams: UpdateItemCommandInput = {
       Key: {
         pk: { S: this.pk },
-        sk: { S: sk },
+        sk: { S: this.buildSk(sk) },
       },
       ReturnValues: 'ALL_NEW',
       TableName: this.tableName,
@@ -156,11 +143,11 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
 
     // To avoid conflict with reserved keywords, refer to attributes with an additional '#' prefix
     for (const key of Object.keys(input)) {
-      if (input[key as keyof Partial<T>] === null) {
+      if (input[key as keyof Partial<Omit<T, keyof SK>>] === null) {
         expressionAttributeNames[`#${key}`] = key;
         expressionAttributeValues[`:${key}`] = { NULL: true };
       } else {
-        const val: string = input[key as keyof Partial<T>] as string;
+        const val: string = input[key as keyof Partial<Omit<T, keyof SK>>] as string;
         expressionAttributeNames[`#${key}`] = key;
         expressionAttributeValues[`:${key}`] = { S: val };
       }
@@ -176,7 +163,7 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
     const updateItemCommand = new UpdateItemCommand(updateParams);
     const updatedRecord = await this.dynamoDbClient.send(updateItemCommand);
 
-    return updatedRecord.Attributes ? this.convertToItem(updatedRecord.Attributes as U) : null;
+    return this.convertToItem(updatedRecord.Attributes as U);
   }
 
   /**
@@ -186,11 +173,11 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
    * @returns {void}
    * @memberof Repository
    */
-  async delete(sk: string): Promise<void> {
+  async delete(sk: T | SK): Promise<void> {
     const deleteCommandParams: DeleteItemCommandInput = {
       Key: {
         pk: { S: this.pk },
-        sk: { S: sk },
+        sk: { S: this.buildSk(sk) },
       },
       TableName: this.tableName,
     };
@@ -199,10 +186,47 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
     return;
   }
 
+  async batchDelete<TItem extends T | SK>(sk: Array<TItem>): Promise<{ unprocessedItems: Array<TItem> }> {
+    if (sk.length === 0) {
+      return { unprocessedItems: [] };
+    }
+
+    const skMap: Record<string, TItem> = {};
+
+    sk.forEach((sk) => {
+      skMap[this.buildSk(sk)] = sk;
+    });
+
+    const result = await this.dynamoDbClient.send(
+      new BatchWriteItemCommand({
+        RequestItems: {
+          [this.tableName]: Object.keys(skMap).map((sk) => ({
+            DeleteRequest: {
+              Key: {
+                pk: { S: this.pk },
+                sk: { S: sk },
+              },
+            },
+          })),
+        },
+      }),
+    );
+
+    return {
+      unprocessedItems:
+        result.UnprocessedItems?.[this.tableName]?.map((item) => {
+          const sk = item.DeleteRequest?.Key?.sk.S;
+
+          // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+          return skMap[sk!];
+        }) || [],
+    };
+  }
+
   /**
    * Convert the DynanomDB attributes of the entity into a DynamoDB item.
    * @abstract
-   * @param {U} attributes The DynamoDB item attributes to convert.
+   * @param {T} item The DynamoDB item attributes to convert.
    * @returns {Record<keyof U, AttributeValue>} The DynamoDB record.
    * @memberof Repository
    */
@@ -216,4 +240,12 @@ export abstract class Repository<T extends BaseItem, U extends BaseDto> {
    * @memberof Repository
    */
   protected abstract convertToItem(record: U): T;
+
+  /**
+   * Converts a sort key to a string for querying/putting in Dynamo.
+   * @param {SK} sk
+   * @return {string}
+   * @protected
+   */
+  protected abstract buildSk(sk: T | SK): string;
 }

package/src/domains/job/repository/JobContextRepository/JobContextRepository.spec.ts

@@ -14,25 +14,26 @@ import {
   ScanCommandOutput,
 } from '@aws-sdk/client-dynamodb';
 import { initIocContainer, iocContainer } from '../../../../mocks/mockIoc';
-import { testLogger } from '../../../../jest-global-setup';
 import {
   createMockClientPutResponse,
   createMockClientScanResponse,
   createMockClientGetResponse,
 } from '../../core/dynamoClient.mock';
-import { PARTITION_KEYS } from '../../constants/Repository.constants';
-import { JobContext } from '../../model/JobContext/JobContext';
-import { createMockJobContext } from '../../model/JobContext/mocks/JobContext.mock';
-import { createMockJobContextDto, createMockJobContextDtos } from '../../model/JobContext/mocks/JobContextDto.mock';
-import { createMockJobContextRecord } from '../../model/JobContext/mocks/JobContextRecord.mock';
-import { createMockCreateJobContextRequest } from '../../model/JobContext/mocks/JobContextRequest.mock';
+import { ENTITY_TYPES } from '../../constants/Repository.constants';
+import {
+  JobContext,
+  createMockJobContext,
+  createMockJobContextDto,
+  createMockJobContextDtos,
+  createMockJobContextRecord,
+  createMockCreateJobContextRequest,
+} from '../../model/JobContext';
 import { JobContextRepository } from './JobContextRepository';
 import * as PaginationUtils from '../../core/pagination/utils/PaginationUtils';
-import { PaginationInfo } from '../../core/pagination/model/PaginationInfo';
-import { PAGINATION_DIRECTION_TYPE_BEFORE } from '../../core/pagination/model/PaginationTypes';
+import { PaginationInfo, PAGINATION_DIRECTION_TYPE_BEFORE } from '../../core/pagination/';
+import { InjectTokens } from '../../constants/InjectTokens';
 
-// Create a new container tsoa can use
-initIocContainer(testLogger);
+initIocContainer();
 
 // Mock DynamoDB Client
 const mockOptions = {} as DynamoDBClientConfig;
@@ -41,20 +42,20 @@ const mockSend = jest.fn();
 mockClient.send = mockSend;
 
 // Inject JobContextRepository with a mocked client
-iocContainer.rebind<DynamoDBClient>(DynamoDBClient).toConstantValue(mockClient);
-
-let jobContextRepository: JobContextRepository;
+iocContainer.rebind(DynamoDBClient).toConstantValue(mockClient);
 
 describe('JobContextRepository', (): void => {
+  let jobContextRepository: JobContextRepository;
+  let dynamoTableName: string;
+  let tenant: string;
+
   beforeAll(() => {
     jobContextRepository = iocContainer.get(JobContextRepository);
+    dynamoTableName = iocContainer.get(InjectTokens.DynamoTableName);
+    tenant = iocContainer.get(InjectTokens.Tenant);
   });
 
   describe('find', () => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return an empty list`, async (): Promise<void> => {
       mockSend.mockImplementationOnce((): ScanCommandOutput => createMockClientScanResponse());
 
@@ -80,8 +81,8 @@ describe('JobContextRepository', (): void => {
     });
 
     it(`should paginate results to return a list of existing job contexts`, async (): Promise<void> => {
-      const mockContextItemsPage1 = createMockJobContextDtos();
-      const mockContextItemsPage2 = createMockJobContextDtos();
+      const mockContextItemsPage1 = createMockJobContextDtos(2);
+      const mockContextItemsPage2 = createMockJobContextDtos(2);
 
       mockSend
         .mockImplementationOnce(
@@ -99,17 +100,15 @@ describe('JobContextRepository', (): void => {
             }),
         );
 
-      const response = await jobContextRepository.list();
+      const response = await jobContextRepository.list({
+        limit: 3,
+      });
       expect(mockSend).toBeCalledTimes(2);
-      expect(response).toHaveLength(mockContextItemsPage1.length + mockContextItemsPage2.length);
+      expect(response).toHaveLength(3);
     });
   });
 
   describe('findOne', () => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return job context by name`, async (): Promise<void> => {
       const mockJobContextItem = createMockJobContextDto({
         environment: { key: 'value' },
@@ -121,7 +120,7 @@ describe('JobContextRepository', (): void => {
           }),
       );
 
-      const response = await jobContextRepository.findOne(mockJobContextItem.sk.S as string);
+      const response = await jobContextRepository.findOne({ contextName: mockJobContextItem.sk.S as string });
       expect(response).toStrictEqual({
         contextName: mockJobContextItem.sk.S as string,
         description: mockJobContextItem.description.S as string,
@@ -139,22 +138,18 @@ describe('JobContextRepository', (): void => {
           }),
       );
 
-      const response = await jobContextRepository.findOne(`null`);
+      const response = await jobContextRepository.findOne({ contextName: `null` });
       expect(mockSend).toBeCalledTimes(1);
       expect(response).toBe(null);
     });
   });
 
   describe('create', () => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should create a new job context record and return job context in the response`, async (): Promise<void> => {
       const newJobContextToCreate = createMockCreateJobContextRequest();
       const mockResponse = createMockClientPutResponse({
         Attributes: createMockJobContextDto({
-          pk: PARTITION_KEYS.context,
+          pk: `${ENTITY_TYPES.context}~${tenant}`,
           sk: newJobContextToCreate.contextName,
         }),
       });
@@ -165,16 +160,12 @@ describe('JobContextRepository', (): void => {
       expect(mockSend).toBeCalledTimes(1);
       expect(response).toStrictEqual({
         ...newJobContextToCreate,
-        type: PARTITION_KEYS.context,
+        type: ENTITY_TYPES.context,
       });
     });
   });
 
   describe('update', () => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should update a job context record`, async (): Promise<void> => {
       const jobContextToUpdate = createMockJobContext();
       const updatedJobContext: JobContext = {
@@ -197,17 +188,16 @@ describe('JobContextRepository', (): void => {
 
       mockSend.mockImplementation((): PutItemCommandOutput => putResponse);
 
-      const response = await jobContextRepository.update(jobContextToUpdate.contextName, updatedJobContextRecord);
+      const response = await jobContextRepository.update(
+        { contextName: jobContextToUpdate.contextName },
+        updatedJobContextRecord,
+      );
       expect(response).toStrictEqual(updatedJobContext);
       expect(mockSend).toBeCalledTimes(1);
     });
   });
 
   describe('delete', () => {
-    beforeEach(() => {
-      jest.resetAllMocks();
-    });
-
     it(`should return undefined when calling the delete operation`, async (): Promise<void> => {
       mockSend.mockImplementationOnce((): DeleteItemCommandOutput => {
         return {
@@ -215,7 +205,7 @@ describe('JobContextRepository', (): void => {
         };
       });
 
-      const response = await jobContextRepository.delete(`name`);
+      const response = await jobContextRepository.delete({ contextName: `name` });
       expect(response).toBe(undefined);
       expect(mockSend).toBeCalledTimes(1);
     });
@@ -228,7 +218,7 @@ describe('JobContextRepository', (): void => {
 
     beforeEach(() => {
       mockExclusiveStartKey = {
-        pk: { S: PARTITION_KEYS.context },
+        pk: { S: ENTITY_TYPES.context },
         sk: { S: 'mockSk' },
       };
       mockPaginationInfo = {
@@ -240,11 +230,11 @@ describe('JobContextRepository', (): void => {
       mockParams = {
         ExclusiveStartKey: mockExclusiveStartKey,
         ExpressionAttributeValues: {
-          ':pk': { S: PARTITION_KEYS.context },
+          ':pk': { S: `${ENTITY_TYPES.context}~${tenant}` },
         },
         KeyConditionExpression: 'pk = :pk',
         ScanIndexForward: false,
-        TableName: process.env.DYNAMODB_TABLE,
+        TableName: dynamoTableName,
       };
 
       jest.spyOn(PaginationUtils, 'decodeTokenString').mockReturnValue(mockExclusiveStartKey);
@@ -271,10 +261,10 @@ describe('JobContextRepository', (): void => {
       expect(result.lastEvaluatedKey).toBe(null);
     });
 
-    it('should return an array of jobs if found', async () => {
+    it('should return an array of contexts if found', async () => {
       const mockJobContextRecords = createMockJobContextDtos();
       const mockLastEvaluatedKey = {
-        pk: { S: PARTITION_KEYS.job },
+        pk: { S: ENTITY_TYPES.job },
         sk: { S: 'mockSk' },
       };