fix tests to use shim.isWrapped instead

jsumners-nr · Nov 7, 2022 · 0e2af07 · 0e2af07
1 parent 0e0024f
commit 0e2af07
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 34 deletions.
diff --git a/lib/instrumentation/koa/tests/unit/koa.tap.js b/lib/instrumentation/koa/tests/unit/koa.tap.js
@@ -14,25 +14,20 @@ tap.test('Koa instrumentation', function (t) {
   const wrapped = ['createContext', 'use', 'emit']
   const notWrapped = ['handleRequest', 'listen', 'toJSON', 'inspect', 'callback', 'onerror']
 
-  // Save the original methods, to compare with wrapped ones below
-  const origKoa = require('koa')
-  const origMethods = Object.fromEntries(
-    wrapped.concat(notWrapped).map((method) => [method, origKoa.prototype[method]])
-  )
-
   const helper = utils.TestAgent.makeInstrumented()
   helper.registerInstrumentation({
     moduleName: 'koa',
     type: 'web-framework',
     onRequire: require('../../lib/instrumentation')
   })
   const Koa = require('koa')
+  const shim = helper.getShim()
 
   wrapped.forEach(function (method) {
-    t.not(Koa.prototype[method], origMethods[method], method + ' is wrapped, as expected')
+    t.ok(shim.isWrapped(Koa.prototype[method]), method + ' is wrapped, as expected')
   })
   notWrapped.forEach(function (method) {
-    t.equal(Koa.prototype[method], origMethods[method], method + ' is not wrapped, as expected')
+    t.not(shim.isWrapped(Koa.prototype[method]), method + ' is not wrapped, as expected')
   })
 
   helper && helper.unload()

diff --git a/lib/instrumentation/koa/tests/unit/route.tap.js b/lib/instrumentation/koa/tests/unit/route.tap.js
@@ -22,10 +22,12 @@ tap.test('koa-route', function (t) {
     onRequire: require('../../lib/route-instrumentation.js')
   })
 
+  const shim = helper.getShim()
+
   t.test('methods', function (t) {
     const route = require('koa-route')
     METHODS.forEach(function checkWrapped(method) {
-      t.ok(route[method].name.startsWith('wrapped'), method + ' should be wrapped')
+      t.ok(shim.isWrapped(route[method]), method + ' should be wrapped')
     })
     t.end()
   })

diff --git a/lib/instrumentation/koa/tests/unit/router.tap.js b/lib/instrumentation/koa/tests/unit/router.tap.js
@@ -25,22 +25,11 @@ const koaRouterMods = ['koa-router', '@koa/router']
 
 koaRouterMods.forEach((koaRouterMod) => {
   tap.test(koaRouterMod, function tests(t) {
-    // Save the original methods, to compare with wrapped ones below
-    const origRouter = require(koaRouterMod)
-    const origMethods = Object.fromEntries(
-      [...WRAPPED_METHODS, ...UNWRAPPED_METHODS].map((method) => [
-        method,
-        origRouter.prototype[method]
-      ])
-    )
-    const origStaticMethods = Object.fromEntries(
-      UNWRAPPED_STATIC_METHODS.map((method) => [method, origRouter[method]])
-    )
-
-    var helper = utils.TestAgent.makeInstrumented()
+    const helper = utils.TestAgent.makeInstrumented()
     t.teardown(function () {
       helper.unload()
     })
+    const shim = helper.getShim()
 
     helper.registerInstrumentation({
       type: 'web-framework',
@@ -52,32 +41,26 @@ koaRouterMods.forEach((koaRouterMod) => {
       var Router = require(koaRouterMod)
       var router = new Router()
       router.param('second', function () {})
-      t.ok(router.params.second.name.startsWith('wrap'), 'param function should be wrapped')
+      t.ok(shim.isWrapped(router.params.second), 'param function should be wrapped')
       t.end()
     })
 
     t.test('methods', function (t) {
       var Router = require(koaRouterMod)
       WRAPPED_METHODS.forEach(function checkWrapped(method) {
-        t.not(
-          Router.prototype[method],
-          origMethods[method],
+        t.ok(
+          shim.isWrapped(Router.prototype[method]),
           method + ' should be a wrapped method on the prototype'
         )
       })
       UNWRAPPED_METHODS.forEach(function checkUnwrapped(method) {
-        t.equal(
-          Router.prototype[method],
-          origMethods[method],
+        t.not(
+          shim.isWrapped(Router.prototype[method]),
           method + ' should be a unwrapped method on the prototype'
         )
       })
       UNWRAPPED_STATIC_METHODS.forEach(function checkUnwrappedStatic(method) {
-        t.equal(
-          Router[method],
-          origStaticMethods[method],
-          method + ' should be an unwrapped static method'
-        )
+        t.not(shim.isWrapped(Router[method]), method + ' should be an unwrapped static method')
       })
       t.end()
     })