@@ -26,6 +26,7 @@ import {
2626 applyInstrumentationPatches ,
2727 AWSXRAY_TRACE_ID_HEADER_CAPITALIZED ,
2828 customExtractor ,
29+ ExtendedAwsLambdaInstrumentation ,
2930 headerGetter ,
3031} from './../../src/patches/instrumentation-patch' ;
3132import * as sinon from 'sinon' ;
@@ -37,6 +38,7 @@ import * as nock from 'nock';
3738import { ReadableSpan , Span as SDKSpan } from '@opentelemetry/sdk-trace-base' ;
3839import { getTestSpans } from '@opentelemetry/contrib-test-utils' ;
3940import { instrumentationConfigs } from '../../src/register' ;
41+ import { LoggerProvider } from '@opentelemetry/api-logs' ;
4042
4143// It is assumed that bedrock.test.ts has already registered the
4244// necessary instrumentations for testing by calling:
@@ -673,6 +675,124 @@ describe('InstrumentationPatchTest', () => {
673675 } ) ;
674676 } ) ;
675677 } ) ;
678+
679+ describe ( 'AwsLambdaInstrumentationPatchTest' , ( ) => {
680+ let awsLambdaInstrumentation : ExtendedAwsLambdaInstrumentation ;
681+ beforeEach ( ( ) => {
682+ const instrumentationsToTest = [
683+ new AwsLambdaInstrumentation ( instrumentationConfigs [ '@opentelemetry/instrumentation-aws-lambda' ] ) ,
684+ ] ;
685+ applyInstrumentationPatches ( instrumentationsToTest ) ;
686+ awsLambdaInstrumentation = instrumentationsToTest [ 0 ] as ExtendedAwsLambdaInstrumentation ;
687+ } ) ;
688+
689+ afterEach ( ( ) => {
690+ sinon . restore ( ) ;
691+ } ) ;
692+
693+ it ( 'Tests setLoggerProvider method' , ( ) => {
694+ const resolvedPromise = Promise . resolve ( ) ;
695+ const mockLoggerProvider = {
696+ getDelegate : ( ) => ( {
697+ forceFlush : ( ) => resolvedPromise ,
698+ } ) ,
699+ getLogger : ( ) => {
700+ return {
701+ emit : ( ) => { } ,
702+ } ;
703+ } ,
704+ } ;
705+
706+ awsLambdaInstrumentation . setLoggerProvider ( mockLoggerProvider as LoggerProvider ) ;
707+ expect ( awsLambdaInstrumentation [ '_logForceFlusher' ] ! ( ) ) . toBe ( resolvedPromise ) ;
708+ } ) ;
709+
710+ it ( 'Tests _logForceFlush with provider that has getDelegate' , ( ) => {
711+ const mockForceFlush = sinon . stub ( ) . resolves ( ) ;
712+ const mockLoggerProvider = {
713+ getDelegate : ( ) => ( {
714+ forceFlush : mockForceFlush ,
715+ } ) ,
716+ getLogger : ( ) => {
717+ return {
718+ emit : ( ) => { } ,
719+ } ;
720+ } ,
721+ } ;
722+
723+ const flusher = awsLambdaInstrumentation [ '_logForceFlush' ] ( mockLoggerProvider as LoggerProvider ) ;
724+ expect ( flusher ) . toBeDefined ( ) ;
725+ flusher ?.( ) ;
726+ expect ( mockForceFlush . called ) . toBeTruthy ( ) ;
727+ } ) ;
728+
729+ it ( 'Tests _logForceFlush with provider that has direct forceFlush' , ( ) => {
730+ const mockForceFlush = sinon . stub ( ) . resolves ( ) ;
731+ const mockLoggerProvider = {
732+ forceFlush : mockForceFlush ,
733+ getLogger : ( ) => {
734+ return {
735+ emit : ( ) => { } ,
736+ } ;
737+ } ,
738+ } ;
739+
740+ const flusher = awsLambdaInstrumentation [ '_logForceFlush' ] ( mockLoggerProvider as LoggerProvider ) ;
741+ expect ( flusher ) . toBeDefined ( ) ;
742+ flusher ?.( ) ;
743+ expect ( mockForceFlush . called ) . toBeTruthy ( ) ;
744+ } ) ;
745+
746+ it ( 'Tests _logForceFlush with undefined provider' , ( ) => {
747+ const flusher = awsLambdaInstrumentation [ '_logForceFlush' ] ( undefined as unknown as LoggerProvider ) ;
748+ expect ( flusher ) . toBeUndefined ( ) ;
749+ } ) ;
750+
751+ it ( 'Tests _endSpan with all flushers' , done => {
752+ const mockSpan : Span = sinon . createStubInstance ( SDKSpan ) ;
753+
754+ // Setup mock flushers
755+ const mockTraceFlush = sinon . stub ( ) . resolves ( ) ;
756+ const mockMetricFlush = sinon . stub ( ) . resolves ( ) ;
757+ const mockLogFlush = sinon . stub ( ) . resolves ( ) ;
758+
759+ awsLambdaInstrumentation [ '_traceForceFlusher' ] = mockTraceFlush ;
760+ awsLambdaInstrumentation [ '_metricForceFlusher' ] = mockMetricFlush ;
761+ awsLambdaInstrumentation [ '_logForceFlusher' ] = mockLogFlush ;
762+
763+ awsLambdaInstrumentation [ '_endSpan' ] ( mockSpan , null , ( ) => {
764+ expect ( mockTraceFlush . called ) . toBeTruthy ( ) ;
765+ expect ( mockMetricFlush . called ) . toBeTruthy ( ) ;
766+ expect ( mockLogFlush . called ) . toBeTruthy ( ) ;
767+ done ( ) ;
768+ } ) ;
769+ } ) ;
770+
771+ it ( 'Tests _endSpan handles missing flushers gracefully' , done => {
772+ const mockSpan : Span = sinon . createStubInstance ( SDKSpan ) ;
773+ const mockDiag = sinon . spy ( diag , 'debug' ) ;
774+ const mockDiagError = sinon . spy ( diag , 'error' ) ;
775+
776+ awsLambdaInstrumentation [ '_endSpan' ] ( mockSpan , null , ( ) => {
777+ expect (
778+ mockDiagError . calledWith (
779+ 'Spans may not be exported for the lambda function because we are not force flushing before callback.'
780+ )
781+ ) . toBeTruthy ( ) ;
782+ expect (
783+ mockDiag . calledWith (
784+ 'Metrics may not be exported for the lambda function because we are not force flushing before callback.'
785+ )
786+ ) . toBeTruthy ( ) ;
787+ expect (
788+ mockDiag . calledWith (
789+ 'Logs may not be exported for the lambda function because we are not force flushing before callback.'
790+ )
791+ ) . toBeTruthy ( ) ;
792+ done ( ) ;
793+ } ) ;
794+ } ) ;
795+ } ) ;
676796} ) ;
677797
678798describe ( 'customExtractor' , ( ) => {
0 commit comments