From d77f7fdc60e9b03e0665ac56b40ed7f0e4789215 Mon Sep 17 00:00:00 2001
From: Jake Ginnivan <jake@ginnivan.net>
Date: Mon, 11 Dec 2023 17:28:40 +0800
Subject: [PATCH] Fixed retry tests

---
 libs/js-sdk/src/utils/retry.spec.ts | 14 ++++++++------
 libs/js-sdk/src/utils/retry.ts      |  8 +++++---
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/libs/js-sdk/src/utils/retry.spec.ts b/libs/js-sdk/src/utils/retry.spec.ts
index 90f7da3..c8a9e64 100644
--- a/libs/js-sdk/src/utils/retry.spec.ts
+++ b/libs/js-sdk/src/utils/retry.spec.ts
@@ -26,7 +26,7 @@ describe('retry function with OpenTelemetry', () => {
     it('should retry the function and succeed', async () => {
         let attempt = 0
         const mockFn = async () => {
-            if (attempt < 2) {
+            if (attempt < 1) {
                 attempt++
                 throw new Error('Temporary failure')
             }
@@ -37,12 +37,14 @@ describe('retry function with OpenTelemetry', () => {
 
         expect(result).toBe('Success')
         const spans = exporter.getFinishedSpans()
-        expect(spans.length).toEqual(5)
+        console.log(spans)
+        expect(spans.length).toEqual(4)
+        expect(spans[0].name).toBe('retry-attempt')
         expect(spans[0].attributes?.retryAttempt).toBe(0)
-        expect(spans[1].name).toEqual('delay')
-        expect(spans[2].attributes?.retryAttempt).toBe(1)
-        expect(spans[3].name).toEqual('delay')
-        expect(spans[4].attributes?.retryAttempt).toBe(2)
+        expect(spans[1].name).toEqual('retry')
+        expect(spans[2].name).toEqual('delay')
+        expect(spans[3].name).toBe('retry-attempt')
+        expect(spans[3].attributes?.retryAttempt).toBe(1)
     })
 
     it('should retry the function and fail', async () => {
diff --git a/libs/js-sdk/src/utils/retry.ts b/libs/js-sdk/src/utils/retry.ts
index df56315..51816d9 100644
--- a/libs/js-sdk/src/utils/retry.ts
+++ b/libs/js-sdk/src/utils/retry.ts
@@ -19,9 +19,11 @@ export async function retry<T>(
         // eslint-disable-next-line no-constant-condition
         while (true) {
             try {
-                return await retryAttemptFn<T>(tracer, retryAttempt, fn).then(
-                    () => span.end(),
-                )
+                return await retryAttemptFn<T>(
+                    tracer,
+                    retryAttempt,
+                    fn,
+                ).finally(() => span.end())
             } catch (error) {
                 const err = resolveError(error)
                 if (cancellationToken?.cancel) {