Feature parity w/ Python #851

Author: dreamorosi

packages/tracing/src/Tracer.ts

@@ -173,12 +173,11 @@ class Tracer implements TracerInterface {
    * @see https://docs.aws.amazon.com/lambda/latest/dg/runtimes-context.html
    */
   public annotateColdStart(): void {
+    if (this.tracingEnabled === true) {
+      this.putAnnotation('ColdStart', Tracer.coldStart);
+    }
     if (Tracer.coldStart === true) {
       Tracer.coldStart = false;
-
-      if (this.tracingEnabled === true) {
-        this.putAnnotation('ColdStart', true);
-      }
     }
   }
 

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

@@ -40,7 +40,7 @@ describe('Class: Tracer', () => {
 
     });
 
-    test('when called multiple times, it annotates the first time and then never again', () => {
+    test('when called multiple times, it annotates true the first time and then false afterwards', () => {
 
       // Prepare
       const tracer: Tracer = new Tracer();
@@ -53,8 +53,13 @@ describe('Class: Tracer', () => {
       tracer.annotateColdStart();
 
       // Assess
-      expect(putAnnotationSpy).toBeCalledTimes(1);
-      expect(putAnnotationSpy).toBeCalledWith('ColdStart', true);
+      expect(putAnnotationSpy).toBeCalledTimes(4);
+      expect(putAnnotationSpy.mock.calls).toEqual([
+        [ 'ColdStart', true ],
+        [ 'ColdStart', false ],
+        [ 'ColdStart', false ],
+        [ 'ColdStart', false ],
+      ]);
 
     });
 
@@ -666,7 +671,7 @@ describe('Class: Tracer', () => {
         .mockImplementation(() => newSubsegmentSecondInvocation);
       setContextMissingStrategy(() => null);
       const captureAsyncFuncSpy = jest.spyOn(tracer.provider, 'captureAsyncFunc');
-      const addAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
+      const putAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
       class Lambda implements LambdaInterface {
 
         @tracer.captureLambdaHanlder()
@@ -687,16 +692,22 @@ describe('Class: Tracer', () => {
       // Assess
       expect(captureAsyncFuncSpy).toHaveBeenCalledTimes(2);
       expect(captureAsyncFuncSpy).toHaveBeenCalledWith('## foo-bar-function', expect.anything());
-      expect(addAnnotationSpy).toHaveBeenCalledTimes(1);
-      expect(addAnnotationSpy).toHaveBeenCalledWith('ColdStart', true);
+      expect(putAnnotationSpy).toHaveBeenCalledTimes(2);
+      expect(putAnnotationSpy.mock.calls).toEqual([
+        [ 'ColdStart', true ],
+        [ 'ColdStart', false ],
+      ]);
       expect(newSubsegmentFirstInvocation).toEqual(expect.objectContaining({
         name: '## foo-bar-function',
         annotations: {
           'ColdStart': true,
         }
       }));
       expect(newSubsegmentSecondInvocation).toEqual(expect.objectContaining({
-        name: '## foo-bar-function'
+        name: '## foo-bar-function',
+        annotations: {
+          'ColdStart': false,
+        }
       }));
 
     });

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

@@ -216,7 +216,7 @@ describe('Middy middlewares', () => {
       .mockImplementationOnce(() => newSubsegmentFirstInvocation)
       .mockImplementation(() => newSubsegmentSecondInvocation);
     setContextMissingStrategy(() => null);
-    const addAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
+    const putAnnotationSpy = jest.spyOn(tracer, 'putAnnotation');
     const lambdaHandler: Handler = async (_event: unknown, _context: Context) => ({
       foo: 'bar'
     });
@@ -232,16 +232,22 @@ describe('Middy middlewares', () => {
     expect(setSegmentSpy).toHaveBeenCalledWith(expect.objectContaining({
       name: '## foo-bar-function',
     }));
-    expect(addAnnotationSpy).toHaveBeenCalledTimes(1);
-    expect(addAnnotationSpy).toHaveBeenCalledWith('ColdStart', true);
+    expect(putAnnotationSpy).toHaveBeenCalledTimes(2);
+    expect(putAnnotationSpy.mock.calls).toEqual([
+      [ 'ColdStart', true ],
+      [ 'ColdStart', false ],
+    ]);
     expect(newSubsegmentFirstInvocation).toEqual(expect.objectContaining({
       name: '## foo-bar-function',
       annotations: {
         'ColdStart': true,
       }
     }));
     expect(newSubsegmentSecondInvocation).toEqual(expect.objectContaining({
-      name: '## foo-bar-function'
+      name: '## foo-bar-function',
+      annotations: {
+        'ColdStart': false,
+      }
     }));
 
   });