Typing & Tests

Author: dreamorosi

packages/tracing/src/Tracer.ts

@@ -25,20 +25,20 @@ class Tracer implements ClassThatTraces {
     this.provider = new ProviderService();
   }
 
-  public captureAWS<T>(aws: T): void | T {
-    if (this.tracingEnabled === false) return;
+  public captureAWS<T>(aws: T): T {
+    if (this.tracingEnabled === false) return aws;
 
     return this.provider.captureAWS(aws);
   }
 
-  public captureAWSClient<T>(service: T): void | T {
-    if (this.tracingEnabled === false) return;
+  public captureAWSClient<T>(service: T): T {
+    if (this.tracingEnabled === false) return service;
 
     return this.provider.captureAWSClient(service);
   }
 
-  public captureAWSv3Client<T>(service: T): void | T {
-    if (this.tracingEnabled === false) return;
+  public captureAWSv3Client<T>(service: T): T {
+    if (this.tracingEnabled === false) return service;
 
     return this.provider.captureAWSv3Client(service);
   }
@@ -59,7 +59,7 @@ class Tracer implements ClassThatTraces {
             result = await originalMethod?.apply(this, [ event, context, callback ]);
             this.addResponseAsMetadata(result, context.functionName);
           } catch (error) {
-            this.addErrorAsMetadata(error);
+            this.addErrorAsMetadata(error as Error);
             // TODO: should this error be thrown?? If thrown we get a ERR_UNHANDLED_REJECTION. If not aren't we are basically catching a Customer error?
             // throw error;
           } finally {
@@ -85,8 +85,13 @@ class Tracer implements ClassThatTraces {
     };
   }
 
-  public getSegment(): Segment | Subsegment | undefined {
-    return this.provider.getSegment();
+  public getSegment(): Segment | Subsegment {
+    const segment = this.provider.getSegment();
+    if (segment === undefined) {
+      throw new Error('Failed to get the current sub/segment from the context.');
+    }
+
+    return segment;
   }
 
   public static isColdStart(): boolean {
@@ -131,7 +136,9 @@ class Tracer implements ClassThatTraces {
 
   private addErrorAsMetadata(error: Error): void {
     const subsegment = this.getSegment();
-    if (this.captureError === false || subsegment === undefined || this.tracingEnabled === false) {
+    if (this.captureError === false) {
+      subsegment.addErrorFlag();
+
       return;
     }
 

packages/tracing/tests/unit/Tracer.test.ts

@@ -52,6 +52,19 @@ describe('Class: Tracer', () => {
       }).toThrow('Failed to get the current sub/segment from the context.');
 
     });
+
+    test('when called outside of a namespace or without parent segment, it throws an error', () => {
+
+      // Prepare
+      const tracer: Tracer = new Tracer();
+      jest.spyOn(tracer.provider, 'getSegment').mockImplementation(() => undefined);
+    
+      // Act / Assess
+      expect(() => {
+        tracer.getSegment();
+      }).toThrow('Failed to get the current sub/segment from the context.');
+
+    });
 
     test('when called within a namespace, it returns the parent segment', () => {
 
@@ -426,6 +439,7 @@ describe('Class: Tracer', () => {
       setContextMissingStrategy(() => null);
       const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
       const addErrorSpy = jest.spyOn(newSubsegment, 'addError');
+      const addErrorFlagSpy = jest.spyOn(newSubsegment, 'addErrorFlag');
       class Lambda implements LambdaInterface {
 
         @tracer.captureLambdaHanlder()
@@ -446,6 +460,7 @@ describe('Class: Tracer', () => {
         name: '## foo-bar-function',
       }));
       expect('cause' in newSubsegment).toBe(false);
+      expect(addErrorFlagSpy).toHaveBeenCalledTimes(1);
       expect(addErrorSpy).toHaveBeenCalledTimes(0);
 
       delete process.env.POWERTOOLS_TRACER_CAPTURE_ERROR;
@@ -494,7 +509,7 @@ describe('Class: Tracer', () => {
       const newSubsegmentSecondInvocation: Segment | Subsegment | undefined = new Subsegment('## foo-bar-function');
       jest.spyOn(tracer.provider, 'getSegment')
         .mockImplementationOnce(() => newSubsegmentFirstInvocation)
-        .mockImplementationOnce(() => newSubsegmentSecondInvocation);
+        .mockImplementation(() => newSubsegmentSecondInvocation);
       setContextMissingStrategy(() => null);
       const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
       const addAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
@@ -668,5 +683,4 @@ describe('Class: Tracer', () => {
     });
 
   });
-
 });