fix: made wrapFunction unique to functions (#6796)

Author: pabloerhard

packages/datadog-instrumentations/src/cucumber.js

@@ -259,7 +259,7 @@ function wrapRun (pl, isLatestVersion, version) {
     testStartCh.runStores(ctx, () => {})
     const promises = {}
     try {
-      this.eventBroadcaster.on('envelope', shimmer.wrapFunction(null, () => async (testCase) => {
+      this.eventBroadcaster.on('envelope', async (testCase) => {
         // Only supported from >=8.0.0
         if (testCase?.testCaseFinished) {
           const { testCaseFinished: { willBeRetried } } = testCase
@@ -289,7 +289,7 @@ function wrapRun (pl, isLatestVersion, version) {
             testStartCh.runStores(newCtx, () => {})
           }
         }
-      }))
+      })
       let promise
 
       testFnCh.runStores(ctx, () => {

packages/datadog-instrumentations/src/mariadb.js

@@ -82,12 +82,7 @@ function createWrapQueryCallback (options) {
 
       const cb = arguments[arguments.length - 1]
       const ctx = { sql, conf: options }
-
-      if (typeof cb !== 'function') {
-        arguments.length += 1
-      }
-
-      arguments[arguments.length - 1] = shimmer.wrapFunction(cb, cb => function (err) {
+      const wrapper = (cb) => function (err) {
         if (err) {
           ctx.error = err
           errorCh.publish(ctx)
@@ -96,7 +91,14 @@ function createWrapQueryCallback (options) {
         return typeof cb === 'function'
           ? finishCh.runStores(ctx, cb, this, ...arguments)
           : finishCh.publish(ctx)
-      })
+      }
+
+      if (typeof cb === 'function') {
+        arguments[arguments.length - 1] = shimmer.wrapFunction(cb, wrapper)
+      } else {
+        arguments.length += 1
+        arguments[arguments.length - 1] = wrapper()
+      }
 
       return startCh.runStores(ctx, query, this, ...arguments)
     }

packages/datadog-shimmer/src/shimmer.js

@@ -73,7 +73,7 @@ function copyObjectProperties (original, wrapped, skipKey) {
  * @returns {Function} The wrapped function.
  */
 function wrapFunction (original, wrapper) {
-  if (typeof original === 'object' && original !== null) return original
+  if (typeof original !== 'function') return original
 
   const wrapped = wrapper(original)
 

packages/datadog-shimmer/test/shimmer.spec.js

@@ -372,10 +372,10 @@ describe('shimmer', () => {
       expect(() => shimmer.wrap(() => {}, () => {})).to.throw()
     })
 
-    it('should work with null instead of function', () => {
+    it('should not work with null instead of function', () => {
       const a = null
       const wrapped = shimmer.wrapFunction(a, x => () => x)
-      expect(wrapped()).to.equal(a)
+      expect(typeof wrapped).to.not.equal('function')
     })
 
     it('should not work with an object', () => {