CLS appears not running inside jest integration tests

[gregoryorton-ws]

Observed behavior

When I curl request to my app that's running outside of test environment, the CLS module seems to work. When I run in my integration test, it does not seem to work.

The biggest issue I have is that the test endpoint I'm using will and should result in the user context name and uuid being written to columns in the database createdBy and createdByName - which works from a normal curl call, but when I run from the test, the values are missing. Can you see any problem with the code? Pulling my hair out as to why it wouldn't be working in this test context...

1. Example curl with "full app"

curl -X POST -H 'x-api-key: time-off-api-key' -H 'Authorization: Bearer ...' http://localhost:3391/v1/test/balance-entry

2. Example test:

beforeAll(async () => {
    fakeUser = {
      id: 1,
      uuid: 'testUserId',
      email: 'test@testUserEmail.com',
      name: 'teatUserName',
      authToken: 'testAuthToken',
    };

    const testingModule = await Test.createTestingModule({
      imports: [AppModule],
    })
      .overrideProvider(PrismaService)
      .useValue(jestPrisma)
      .compile();

    app = createNestApplication(testingModule);
    await app.init();
    dbService = testingModule.get(PrismaService);
    usersService = testingModule.get(UsersService);
  });

it('should get user from context and save it in relevant models when creating records', async () => {
    jest
      .spyOn(usersService, 'getCurrentUserFromAuthorizationToken')
      .mockResolvedValue(fakeUser);

    await request(app.getHttpServer())
      .post(`/v1/test/balance-entry`)
      .set({ authorization: 'Bearer fakeToken' })
      .expect(201);

    const timeOffHistoryWithUserContext =
      await dbService.client.timeOffHistory.findMany();

    expect(timeOffHistoryWithUserContext).toEqual(
      expect.arrayContaining([
        expect.objectContaining({
          createdBy: fakeUser.uuid,
          createdByName: fakeUser.name,
        }),
      ])
    );
  });

result from test:

Expected: ArrayContaining [ObjectContaining {"createdBy": "testUserId", "createdByName": "teatUserName"}]
    Received: [{"adjustmentType": "adjustment", "balance": "10", "createdAt": 2024-04-27T10:07:21.508Z, "createdBy": null, "createdByName": null, "hours": "10", "id": "...", "name": null, "notes": null, "policyId": null, "teamMemberId": "testTeamMemberId", "type": "...", "updatedAt": 2024-04-27T10:07:21.508Z}]

Expected behaviour

It operates the same in tests as well as running the full

additional code

// app.module.ts
@Module({
  imports: [
    ClsModule.forRoot({
      global: true,
      middleware: { mount: true },
    }),
    ConfigModule.forRoot({
      isGlobal: true,
      ignoreEnvFile: true,
      load: [applicationConfiguration, bullmqConfig],
    }),
    BullModule.forRootAsync({
      imports: [ConfigModule],
      useFactory: async (configService: ConfigService) =>
        configService.get('bullmq', { connection: {} }),
      inject: [ConfigService],
    }),
    CacheModule.register({
      isGlobal: true,
      ttl: 1000 * 60 * 60, // 1 hour
    }),
    WinstonModule.forRoot(currentWinstonLoggerConfig),
    DatabaseModule,
    ConditionalModule.registerWhen(
      TestingModule,
      (env: NodeJS.ProcessEnv) =>
        env.APP_STAGE === AppStage.Local || env.NODE_ENV === 'test',
      {
        timeout: 150,
      }
    ),
    ConditionalModule.registerWhen(
      KafkaModule,
      (env) => env.KAFKA_ENABLED === 'true',
      {
        timeout: 150,
      }
    ),
    WSCoreModule,
    PayrollModule,
    // Domain Module below:
    UsersModule,
    TimeOffHistoryModule,
    PoliciesModule,
    TeamMembersModule,
    AccrualModule,
    BalancesModule,
  ],
  controllers: [AppController],
  providers: [AppService],
})
export class AppModule {}

// prisma.service.ts
import {
  Injectable,
  Logger,
  OnModuleDestroy,
  OnModuleInit,
} from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { PrismaClient, Prisma } from '@prisma-clients/time-off';
import { ClsService } from 'nestjs-cls';

import { TimeOffService } from '../configuration/interfaces/time-off-service';
import { AppStage } from '../configuration/interfaces/application-config';
import { CustomTimeOffException } from '../common/exceptions';
import { User } from '../users/interfaces/user';

import {
  ConstraintValidationError,
  DatabaseError,
  RecordNotFoundError,
  ValueErrorException,
} from './exceptions/exceptions';

@Injectable()
export class PrismaService implements OnModuleInit, OnModuleDestroy {
  private readonly logger = new Logger(PrismaService.name);
  private readonly clsService: ClsService;

  // TODO: add more error codes to this map to catch more prisma errors and wrap them in custom exceptions
  private readonly prismaCodeToExceptionMap: Record<string, unknown> = {
    P2000: ValueErrorException,
    P2002: ConstraintValidationError,
    P2025: RecordNotFoundError,
  };

  /**
   * Expose the client as a property to enable extending (middleware with $use has been deprecated)
   * https://www.prisma.io/docs/orm/prisma-client/client-extensions/query
   */
  client;

  constructor(
    private readonly configService: ConfigService,
    private readonly cls: ClsService
  ) {
    const _client = new PrismaClient({
      log: [
        {
          emit: 'event',
          level: 'query',
        },
        {
          emit: 'event',
          level: 'error',
        },
        {
          emit: 'stdout',
          level: 'info',
        },
        {
          emit: 'stdout',
          level: 'warn',
        },
      ],
    });

    const appStage = configService.get<TimeOffService['appStage']>(
      'timeOffService.appStage'
    );

    const logger = this.logger;

    _client.$on('query', (e) => {
      if (appStage === AppStage.Local) {
        logger.log(
          `Query: ${e.query} |  Duration: ${e.duration} }`,
          'PrismaService'
        );
      }
    });

    this.client = _client
      .$extends(this.onCreateExtension())
      .$extends(this.onErrorExtension());

    this.clsService = cls;
  }

  private onErrorExtension() {
    const logger = this.logger;
    const prismaCodeToExceptionMap = this.prismaCodeToExceptionMap;
    return Prisma.defineExtension({
      name: 'Error handling',
      query: {
        $allModels: {
          async $allOperations({ query, args }) {
            try {
              return await query(args);
            } catch (error) {
              if (error instanceof Prisma.PrismaClientKnownRequestError) {
                logger.error(error.message, error.stack, 'PrismaService');
                const exception = prismaCodeToExceptionMap[error.code];
                if (exception) {
                  const Exception = exception as typeof CustomTimeOffException;
                  throw new Exception({
                    messages: [error.message],
                    cause: error,
                  });
                }

                throw new DatabaseError({
                  messages: [error.message],
                  cause: error,
                });
              }
              logger.error('Unhandled database error', error, 'PrismaService');
              throw new DatabaseError({
                messages: [
                  "We're sorry, something went wrong. Please try again later.",
                ],
              });
            }
          },
        },
      },
    });
  }

  private onCreateExtension() {
    return Prisma.defineExtension({
      name: 'Auto created audit fields',
      query: {
        policy: {
          create: async ({ query, args, model }) => {
            const user = this.clsService.get('user') as User | undefined;

            if (model === 'Policy' && user) {
              args.data.createdBy = user.uuid ?? null;
              args.data.createdByName = user.name ?? null;
            }
            return query(args);
          },
          createMany: async ({ query, args, model }) => {
            const user = this.clsService.get('user') as User | undefined;

            if (model === 'Policy' && user) {
              if (Array.isArray(args.data)) {
                args.data.forEach((data) => {
                  data.createdBy = user.uuid ?? null;
                  data.createdByName = user.name ?? null;
                });
              } else {
                args.data.createdBy = user.uuid ?? null;
                args.data.createdByName = user.name ?? null;
              }
            }
            return query(args);
          },
        },
        timeOffHistory: {
          create: async ({ query, args, model }) => {
            const user = this.clsService.get('user') as User | undefined;

            if (model === 'TimeOffHistory' && user) {
              args.data.createdBy = user.uuid ?? null;
              args.data.createdByName = user.name ?? null;
            }
            return query(args);
          },
          createMany: async ({ query, args, model }) => {
            const user = this.clsService.get('user') as User | undefined;

            if (model === 'TimeOffHistory' && user) {
              if (Array.isArray(args.data)) {
                args.data.forEach((data) => {
                  data.createdBy = user.uuid ?? null;
                  data.createdByName = user.name ?? null;
                });
              } else {
                args.data.createdBy = user.uuid ?? null;
                args.data.createdByName = user.name ?? null;
              }
            }
            return query(args);
          },
        },
      },
    });
  }

  async onModuleInit() {
    await this.client.$connect();
  }

  async onModuleDestroy() {
    await this.client.$disconnect();
  }
}

@Injectable()
export class AuthorizationTokenGuard implements CanActivate {
  constructor(
    private readonly usersService: UsersService,
    @Inject(CACHE_MANAGER)
    private readonly cacheManager: Cache,
    private readonly cls: ClsService
  ) {}

  async canActivate(context: ExecutionContext): Promise<boolean> {
    const req = context.switchToHttp().getRequest();
    const authorizationHeaderValue: string = req.headers.authorization;

    if (!authorizationHeaderValue) {
      throw new UnauthorizedException('Missing authorization header');
    }

    const authorizationToken = authorizationHeaderValue.replace(/Bearer\s/, '');

    try {
      const user = await this.cacheManager.wrap(authorizationToken, () =>
        this.usersService.getCurrentUserFromAuthorizationToken(
          authorizationToken
        )
      );

      if (!user) {
        throw new UnauthorizedException('Not authenticated');
      }

      this.cls.set('user', user);

      req.user = user; // set user in context

      return true;
    } catch (e) {
      let message = 'Unknown authentication error';
      if (
        axios.isAxiosError(e) &&
        e.response?.data &&
        e.response.data.message
      ) {
        message = e.response.data.message;
      }
      Logger.error(
        `Authentication with core failed: ${message}`,
        'AuthorizationTokenGuard'
      );
      throw new UnauthorizedException('Not authenticated');
    }
  }
}

@Post('/balance-entry')
  @UseGuards(AuthorizationTokenGuard)
  async createTimeOffHistoryEntry() {
    return this.balanceService.insertJournalsAndUpdateBalances({
      journals: [
        {
          teamMemberId: 'testTeamMemberId',
          type: PolicyType.PTO_Vacation,
          adjustmentType: TimeOffHistoryAdjustmentType.Adjustment,
          hours: 10,
          balance: 10,
        },
      ],
    });
  }

[gregoryorton-ws]

ignore this. It's jestPrisma causing the issue.

seems like it's not compatible with my prisma service extensions

[Papooch]

Thank you for the report and identifying the real issue. I'll go ahead and close this now.