From 57a4dfb77c0408cbd81291c71db770005a0f2b5a Mon Sep 17 00:00:00 2001
From: Maurice Rickard <maurice@mauricerickard.com>
Date: Mon, 8 Jul 2024 17:41:10 -0400
Subject: [PATCH] fix: Refactored benchmark tests to complete async functions
 (#2334)

Signed-off-by: mrickard <maurice@mauricerickard.com>
---
 .../datastore-shim/test-datastore.js          | 15 ++++--
 test/benchmark/promises/shared.js             | 54 ++++++++++++++++---
 test/lib/benchmark.js                         |  5 +-
 3 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/test/benchmark/datastore-shim/test-datastore.js b/test/benchmark/datastore-shim/test-datastore.js
index fc5ae81f23..011f04b1ba 100644
--- a/test/benchmark/datastore-shim/test-datastore.js
+++ b/test/benchmark/datastore-shim/test-datastore.js
@@ -7,15 +7,24 @@
 
 class TestDatastore {
   testOp(cb) {
-    setImmediate(cb)
+    if (typeof cb === 'function') {
+      return setImmediate(cb)
+    }
+    return cb || 'testOp'
   }
 
   testQuery(query, cb) {
-    setImmediate(cb)
+    if (typeof cb === 'function') {
+      return setImmediate(cb)
+    }
+    return cb || 'testQuery'
   }
 
   testBatch(query, cb) {
-    setImmediate(cb)
+    if (typeof cb === 'function') {
+      return setImmediate(cb)
+    }
+    return cb || 'testBatch'
   }
 }
 
diff --git a/test/benchmark/promises/shared.js b/test/benchmark/promises/shared.js
index 0b4b46ae5b..8e4115d654 100644
--- a/test/benchmark/promises/shared.js
+++ b/test/benchmark/promises/shared.js
@@ -27,7 +27,12 @@ const tests = [
         promises.push(prom.then(function first() {}))
         promises.push(prom.then(function second() {}))
       }
-      Promise.all(promises).then(cb)
+      Promise.all(promises).then(() => {
+        if (typeof cb === 'function') {
+          return cb()
+        }
+        return cb || true
+      })
     }
   },
 
@@ -37,7 +42,12 @@ const tests = [
       for (let i = 0; i < NUM_PROMISES; ++i) {
         prom = prom.then(function () {})
       }
-      prom.then(cb)
+      prom.then(() => {
+        if (typeof cb === 'function') {
+          return cb()
+        }
+        return cb || true
+      })
     }
   },
 
@@ -47,7 +57,12 @@ const tests = [
       for (let i = 0; i < NUM_PROMISES / 2; ++i) {
         prom = prom.then(function () {}).catch(function () {})
       }
-      prom.then(cb)
+      prom.then(() => {
+        if (typeof cb === 'function') {
+          return cb()
+        }
+        return cb || true
+      })
     }
   },
 
@@ -57,7 +72,14 @@ const tests = [
       for (let i = 0; i < NUM_PROMISES - 1; ++i) {
         prom = prom.then(function () {})
       }
-      prom.catch(function () {}).then(cb)
+      prom
+        .catch(function () {})
+        .then(() => {
+          if (typeof cb === 'function') {
+            return cb()
+          }
+          return cb || true
+        })
     }
   },
 
@@ -69,7 +91,10 @@ const tests = [
           res()
         })
       }
-      cb()
+      if (typeof cb === 'function') {
+        return cb()
+      }
+      return cb || true
     }
   },
 
@@ -87,7 +112,12 @@ const tests = [
           })
         )
       }
-      Promise.all(promises).then(cb)
+      Promise.all(promises).then(() => {
+        if (typeof cb === 'function') {
+          return cb()
+        }
+        return cb || true
+      })
     }
   },
 
@@ -101,7 +131,12 @@ const tests = [
           })
         })
       }
-      prom.then(cb)
+      prom.then(() => {
+        if (typeof cb === 'function') {
+          return cb()
+        }
+        return cb || true
+      })
     }
   },
 
@@ -110,7 +145,10 @@ const tests = [
       new Promise(function () {
         throw new Error('Whoops!')
       }).catch(() => {})
-      cb()
+      if (typeof cb === 'function') {
+        return cb()
+      }
+      return cb || true
     }
   }
 ]
diff --git a/test/lib/benchmark.js b/test/lib/benchmark.js
index 4c6e69e544..b8bfeeb615 100644
--- a/test/lib/benchmark.js
+++ b/test/lib/benchmark.js
@@ -73,9 +73,6 @@ class Benchmark {
       const prevCpu = process.cpuUsage()
       const testFn = test.fn
 
-      if (test.async) {
-        return testFn(agent, () => after(test, next, executeCb, prevCpu))
-      }
       await testFn(agent)
       return after(test, next, executeCb, prevCpu)
     }
@@ -108,7 +105,7 @@ class Benchmark {
       if (idx >= suite.tests.length) {
         return true
       }
-      return initiator(initiator, suite.tests[idx], idx)
+      return await initiator(initiator, suite.tests[idx], idx)
     }
 
     const afterTestRuns = (initiator, test, samples, idx) => {