diff --git a/src/middleware/request-id.middleware.spec.ts b/src/middleware/request-id.middleware.spec.ts new file mode 100644 index 0000000..a1cd8bb --- /dev/null +++ b/src/middleware/request-id.middleware.spec.ts @@ -0,0 +1,108 @@ +import { describe, it, expect } from 'vitest'; +import { Request, Response } from 'express'; +import { RequestIdMiddleware } from './request-id.middleware'; +import { TraceContext } from '../common/context/trace.context'; + +describe('RequestIdMiddleware', () => { + const middleware = new RequestIdMiddleware(); + + it('should generate a request ID when x-request-id header is absent', () => { + const req = { headers: {} } as Request; + const res = { + setHeader: (key: string, value: string) => { + if (key === 'x-request-id') { + expect(value).toBeDefined(); + expect(typeof value).toBe('string'); + } + }, + } as unknown as Response; + const next = () => {}; + + middleware.use(req as Request, res as Response, next); + + }); + + it('should reuse existing x-request-id header when present', () => { + const existingId = 'existing-request-123'; + const req = { + headers: { 'x-request-id': existingId }, + } as unknown as Request; + const res = { + setHeader: () => {}, + } as unknown as Response; + const next = () => {}; + + middleware.use(req as Request, res as Response, next); + }); + + it('should set x-correlation-id to the request ID when no correlation ID header is present', () => { + const req = { + headers: {}, + } as unknown as Request; + const res = { + setHeader: () => {}, + } as unknown as Response; + const next = () => {}; + + middleware.use(req as Request, res as Response, next); + }); + + it('should propagate correlation ID into AsyncLocalStorage', () => { + const existingId = 'custom-correlation-456'; + const req = { + headers: { 'x-request-id': existingId, 'x-correlation-id': existingId }, + } as unknown as Request; + const res = { + setHeader: () => {}, + } as unknown as Response; + const next = () => {}; + + middleware.use(req as Request, res as Response, next); + + const storedTrace = TraceContext.get(); + expect(storedTrace?.traceId).toBe(existingId); + }); + + it('should propagate generated request ID into AsyncLocalStorage when no correlation header', () => { + const req = { + headers: {}, + } as unknown as Request; + const res = { + setHeader: () => {}, + } as unknown as Response; + const next = () => {}; + + middleware.use(req as Request, res as Response, next); + + const storedTrace = TraceContext.get(); + expect(storedTrace).toBeDefined(); + expect(typeof storedTrace.traceId).toBe('string'); + }); + + it('should propagate correlation ID through nested TraceContext.run calls', () => { + let capturedTraceId: string | undefined; + + const nestedFn = () => { + capturedTraceId = TraceContext.getTraceId(); + return 'result'; + }; + + const req = { + headers: { 'x-request-id': 'nested-test-789' }, + } as unknown as Request; + const res = { + setHeader: () => {}, + } as unknown as Response; + const next = () => {}; + + TraceContext.run( + { traceId: 'nested-test-789' }, + () => { + middleware.use(req as Request, res as Response, next); + nestedFn(); + }, + ); + + expect(capturedTraceId).toBe('nested-test-789'); + }); +}); \ No newline at end of file diff --git a/src/middleware/request-id.middleware.ts b/src/middleware/request-id.middleware.ts index 8b6687c..0d58f28 100644 --- a/src/middleware/request-id.middleware.ts +++ b/src/middleware/request-id.middleware.ts @@ -1,6 +1,7 @@ import { Injectable, NestMiddleware } from '@nestjs/common'; import { NextFunction, Request, Response } from 'express'; import { v7 as uuidv7 } from 'uuid'; +import { TraceContext, TraceContextData } from '../common/context/trace.context'; import { CORRELATION_ID_HEADER, REQUEST_ID_HEADER, @@ -18,10 +19,19 @@ export class RequestIdMiddleware implements NestMiddleware { const requestId = existing && existing.length > 0 ? existing : `req_${uuidv7()}`; req.headers[REQUEST_ID_HEADER] = requestId; - const correlation = req.headers[CORRELATION_ID_HEADER] as string | undefined; - req.headers[CORRELATION_ID_HEADER] = correlation && correlation.length > 0 ? correlation : requestId; + const correlation = + req.headers[CORRELATION_ID_HEADER] as string | undefined; + const correlationId = + correlation && correlation.length > 0 ? correlation : requestId; + req.headers[CORRELATION_ID_HEADER] = correlationId; - res.setHeader(REQUEST_ID_HEADER, requestId); - next(); + const traceData: TraceContextData = { + traceId: correlationId, + }; + + TraceContext.run(traceData, () => { + res.setHeader(REQUEST_ID_HEADER, requestId); + next(); + }); } }