@@ -958,6 +958,56 @@ describe('Class: Tracer', () => {
958958 expect ( await handler ( { } , context , ( ) => console . log ( 'Lambda invoked!' ) ) ) . toEqual ( 'memberVariable:someValue' ) ;
959959
960960 } ) ;
961+
962+ test ( 'when used as decorator on an async method, the method is awaited correctly' , async ( ) => {
963+
964+ // Prepare
965+ const tracer : Tracer = new Tracer ( ) ;
966+ const newSubsegment : Segment | Subsegment | undefined = new Subsegment ( '### dummyMethod' ) ;
967+
968+ jest . spyOn ( tracer . provider , 'getSegment' )
969+ . mockImplementation ( ( ) => newSubsegment ) ;
970+ setContextMissingStrategy ( ( ) => null ) ;
971+ const subsegmentCloseSpy = jest . spyOn ( newSubsegment , 'close' ) . mockImplementation ( ) ;
972+ createCaptureAsyncFuncMock ( tracer . provider , newSubsegment ) ;
973+
974+ class Lambda implements LambdaInterface {
975+ public async dummyMethod ( ) : Promise < void > {
976+ return ;
977+ }
978+
979+ @tracer . captureLambdaHandler ( )
980+ // eslint-disable-next-line @typescript-eslint/ban-ts-comment
981+ // @ts -ignore
982+ public async handler < TEvent , TResult > ( _event : TEvent , _context : Context , _callback : Callback < TResult > ) : void | Promise < void > {
983+ await this . dummyMethod ( ) ;
984+ this . otherDummyMethod ( ) ;
985+
986+ return ;
987+ }
988+
989+ public otherDummyMethod ( ) : void {
990+ return ;
991+ }
992+
993+ }
994+
995+ // Act
996+ const lambda = new Lambda ( ) ;
997+ const otherDummyMethodSpy = jest . spyOn ( lambda , 'otherDummyMethod' ) . mockImplementation ( ) ;
998+ const handler = lambda . handler . bind ( lambda ) ;
999+ await handler ( { } , context , ( ) => console . log ( 'Lambda invoked!' ) ) ;
1000+
1001+ // Assess
1002+ // Here we assert that the otherDummyMethodSpy method is called before the cleanup logic (inside the finally of decorator)
1003+ // that should always be called after the handler has returned. If otherDummyMethodSpy is called after it means the
1004+ // decorator is NOT awaiting the handler which would cause the test to fail.
1005+ const dummyCallOrder = subsegmentCloseSpy . mock . invocationCallOrder [ 0 ] ;
1006+ const otherDummyCallOrder = otherDummyMethodSpy . mock . invocationCallOrder [ 0 ] ;
1007+ expect ( otherDummyCallOrder ) . toBeLessThan ( dummyCallOrder ) ;
1008+
1009+ } ) ;
1010+
9611011 } ) ;
9621012
9631013 describe ( 'Method: captureMethod' , ( ) => {
@@ -1241,6 +1291,53 @@ describe('Class: Tracer', () => {
12411291
12421292 } ) ;
12431293
1294+ test ( 'when used as decorator on an async method, the method is awaited correctly' , async ( ) => {
1295+
1296+ // Prepare
1297+ const tracer : Tracer = new Tracer ( ) ;
1298+ const newSubsegment : Segment | Subsegment | undefined = new Subsegment ( '### dummyMethod' ) ;
1299+
1300+ jest . spyOn ( tracer . provider , 'getSegment' )
1301+ . mockImplementation ( ( ) => newSubsegment ) ;
1302+ setContextMissingStrategy ( ( ) => null ) ;
1303+ const subsegmentCloseSpy = jest . spyOn ( newSubsegment , 'close' ) . mockImplementation ( ) ;
1304+ createCaptureAsyncFuncMock ( tracer . provider , newSubsegment ) ;
1305+
1306+ class Lambda implements LambdaInterface {
1307+ @tracer . captureMethod ( )
1308+ // eslint-disable-next-line @typescript-eslint/ban-ts-comment
1309+ // @ts -ignore
1310+ public async dummyMethod ( ) : Promise < void > {
1311+ return ;
1312+ }
1313+
1314+ public async handler < TEvent , TResult > ( _event : TEvent , _context : Context , _callback : Callback < TResult > ) : Promise < void > {
1315+ await this . dummyMethod ( ) ;
1316+ this . otherDummyMethod ( ) ;
1317+
1318+ return ;
1319+ }
1320+
1321+ public otherDummyMethod ( ) : void {
1322+ return ;
1323+ }
1324+ }
1325+
1326+ // Act
1327+ const lambda = new Lambda ( ) ;
1328+ const otherDummyMethodSpy = jest . spyOn ( lambda , 'otherDummyMethod' ) . mockImplementation ( ) ;
1329+ const handler = lambda . handler . bind ( lambda ) ;
1330+ await handler ( { } , context , ( ) => console . log ( 'Lambda invoked!' ) ) ;
1331+
1332+ // Here we assert that the subsegment.close() (inside the finally of decorator) is called before the other otherDummyMethodSpy method
1333+ // that should always be called after the handler has returned. If subsegment.close() is called after it means the
1334+ // decorator is NOT awaiting the method which would cause the test to fail.
1335+ const dummyCallOrder = subsegmentCloseSpy . mock . invocationCallOrder [ 0 ] ;
1336+ const otherDummyCallOrder = otherDummyMethodSpy . mock . invocationCallOrder [ 0 ] ;
1337+ expect ( dummyCallOrder ) . toBeLessThan ( otherDummyCallOrder ) ;
1338+
1339+ } ) ;
1340+
12441341 test ( 'when used as decorator together with another external decorator, the method name is detected properly' , async ( ) => {
12451342
12461343 // Prepare
0 commit comments