@@ -6,11 +6,13 @@ import {
66 diag ,
77 Context as OtelContext ,
88 trace ,
9+ context ,
910 propagation ,
1011 Span ,
1112 Tracer ,
1213 AttributeValue ,
1314 TextMapSetter ,
15+ INVALID_SPAN_CONTEXT ,
1416} from '@opentelemetry/api' ;
1517import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node' ;
1618import { Instrumentation } from '@opentelemetry/instrumentation' ;
@@ -19,7 +21,12 @@ import { AwsLambdaInstrumentation, AwsLambdaInstrumentationConfig } from '@opent
1921import { expect } from 'expect' ;
2022import { AWS_ATTRIBUTE_KEYS } from '../../src/aws-attribute-keys' ;
2123import { RequestMetadata , ServiceExtension } from '../../src/third-party/otel/aws/services/ServiceExtension' ;
22- import { applyInstrumentationPatches , customExtractor , headerGetter } from './../../src/patches/instrumentation-patch' ;
24+ import {
25+ applyInstrumentationPatches ,
26+ AWSXRAY_TRACE_ID_HEADER_CAPITALIZED ,
27+ customExtractor ,
28+ headerGetter ,
29+ } from './../../src/patches/instrumentation-patch' ;
2330import * as sinon from 'sinon' ;
2431import { AWSXRAY_TRACE_ID_HEADER , AWSXRayPropagator } from '@opentelemetry/propagator-aws-xray' ;
2532import { Context } from 'aws-lambda' ;
@@ -535,34 +542,72 @@ describe('InstrumentationPatchTest', () => {
535542 let lambda : Lambda ;
536543 const region = 'us-east-1' ;
537544
538- it ( 'overridden _getV3SmithyClientSendPatch updates MiddlewareStack' , async ( ) => {
539- const mockedMiddlewareStackInternal : any = [ ] ;
540- const mockedMiddlewareStack = {
541- add : ( arg1 : any , arg2 : any ) => mockedMiddlewareStackInternal . push ( [ arg1 , arg2 ] ) ,
542- } ;
543- const send = extractAwsSdkInstrumentation ( PATCHED_INSTRUMENTATIONS )
544- [ '_getV3SmithyClientSendPatch' ] ( ( ...args : unknown [ ] ) => Promise . resolve ( ) )
545- . bind ( { middlewareStack : mockedMiddlewareStack } ) ;
546- sinon
547- . stub ( AWSXRayPropagator . prototype , 'inject' )
548- . callsFake ( ( context : OtelContext , carrier : unknown , setter : TextMapSetter ) => {
549- ( carrier as any ) [ 'isCarrierModified' ] = 'carrierIsModified' ;
550- } ) ;
545+ describe ( 'overridden _getV3SmithyClientSendPatch updates MiddlewareStack' , async ( ) => {
546+ let mockedMiddlewareStackInternal : any ;
547+ let mockedMiddlewareStack ;
548+ let middlewareArgsHeader : any ;
549+ const testXrayTraceHeader = 'test-xray-trace-header' ;
550+
551+ beforeEach ( async ( ) => {
552+ // Clear environment variables before each test
553+ mockedMiddlewareStackInternal = [ ] ;
554+ mockedMiddlewareStack = {
555+ add : ( arg1 : any , arg2 : any ) => mockedMiddlewareStackInternal . push ( [ arg1 , arg2 ] ) ,
556+ } ;
557+ const send = extractAwsSdkInstrumentation ( PATCHED_INSTRUMENTATIONS )
558+ [ '_getV3SmithyClientSendPatch' ] ( ( ...args : unknown [ ] ) => Promise . resolve ( ) )
559+ . bind ( { middlewareStack : mockedMiddlewareStack } ) ;
560+
561+ middlewareArgsHeader = {
562+ request : {
563+ headers : { } ,
564+ } ,
565+ } ;
566+
567+ // eslint-disable-next-line @typescript-eslint/ban-ts-comment
568+ // @ts -ignore
569+ await send ( { } , null ) ;
570+ } ) ;
551571
552- // eslint-disable-next-line @typescript-eslint/ban-ts-comment
553- // @ts -ignore
554- await send ( { } , null ) ;
572+ afterEach ( ( ) => {
573+ sinon . restore ( ) ;
574+ } ) ;
555575
556- const middlewareArgs : any = {
557- request : {
558- headers : { } ,
559- } ,
560- } ;
561- await mockedMiddlewareStackInternal [ 0 ] [ 0 ] ( ( arg : any ) => Promise . resolve ( ) , null ) ( middlewareArgs ) ;
576+ it ( 'Updates trace header casing when AWSXRayPropagator injects trace header successfully' , async ( ) => {
577+ sinon
578+ . stub ( AWSXRayPropagator . prototype , 'inject' )
579+ . callsFake ( ( context : OtelContext , carrier : unknown , setter : TextMapSetter ) => {
580+ ( carrier as any ) [ 'isCarrierModified' ] = 'carrierIsModified' ;
581+ ( carrier as any ) [ AWSXRAY_TRACE_ID_HEADER ] = testXrayTraceHeader ;
582+ } ) ;
583+ await mockedMiddlewareStackInternal [ 0 ] [ 0 ] ( ( arg : any ) => Promise . resolve ( ) , null ) ( middlewareArgsHeader ) ;
584+
585+ expect ( middlewareArgsHeader . request . headers [ 'isCarrierModified' ] ) . toEqual ( 'carrierIsModified' ) ;
586+ expect ( middlewareArgsHeader . request . headers ) . not . toHaveProperty ( AWSXRAY_TRACE_ID_HEADER ) ;
587+ expect ( middlewareArgsHeader . request . headers ) . toHaveProperty ( AWSXRAY_TRACE_ID_HEADER_CAPITALIZED ) ;
588+ expect ( middlewareArgsHeader . request . headers [ AWSXRAY_TRACE_ID_HEADER_CAPITALIZED ] ) . toEqual ( testXrayTraceHeader ) ;
589+
590+ expect ( mockedMiddlewareStackInternal [ 0 ] [ 1 ] . name ) . toEqual ( '_adotInjectXrayContextMiddleware' ) ;
591+ } ) ;
562592
563- sinon . restore ( ) ;
564- expect ( middlewareArgs . request . headers [ 'isCarrierModified' ] ) . toEqual ( 'carrierIsModified' ) ;
565- expect ( mockedMiddlewareStackInternal [ 0 ] [ 1 ] . name ) . toEqual ( '_adotInjectXrayContextMiddleware' ) ;
593+ it ( 'Does not set trace header when AWSXRayPropagator does not inject trace header' , async ( ) => {
594+ const invalidContext : OtelContext = {
595+ getValue : ( key : symbol ) => ( {
596+ spanContext : ( ) => INVALID_SPAN_CONTEXT ,
597+ } ) ,
598+ setValue : ( key : symbol , value : unknown ) => invalidContext ,
599+ deleteValue : ( key : symbol ) => invalidContext ,
600+ } ;
601+
602+ sinon . stub ( context , 'active' ) . returns ( invalidContext ) ;
603+
604+ await mockedMiddlewareStackInternal [ 0 ] [ 0 ] ( ( arg : any ) => Promise . resolve ( ) , null ) ( middlewareArgsHeader ) ;
605+
606+ expect ( middlewareArgsHeader . request . headers ) . not . toHaveProperty ( AWSXRAY_TRACE_ID_HEADER ) ;
607+ expect ( middlewareArgsHeader . request . headers ) . not . toHaveProperty ( AWSXRAY_TRACE_ID_HEADER_CAPITALIZED ) ;
608+
609+ expect ( mockedMiddlewareStackInternal [ 0 ] [ 1 ] . name ) . toEqual ( '_adotInjectXrayContextMiddleware' ) ;
610+ } ) ;
566611 } ) ;
567612
568613 it ( 'injects trace context header into request via propagator' , async ( ) => {
0 commit comments