From 84d3d9c07e4f4a91dec9eb1d462e05894cbd0757 Mon Sep 17 00:00:00 2001
From: Yongseok Kang <feelform@gmail.com>
Date: Fri, 11 Feb 2022 15:02:19 +0900
Subject: [PATCH] [#104] koa support code level visibility

---
 lib/instrumentation/module/koa-router.js | 15 +++++++++++++--
 test/instrumentation/module/koa.test.js  |  7 +++----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/lib/instrumentation/module/koa-router.js b/lib/instrumentation/module/koa-router.js
index f9b98e96..c57da0c7 100644
--- a/lib/instrumentation/module/koa-router.js
+++ b/lib/instrumentation/module/koa-router.js
@@ -52,14 +52,14 @@ module.exports = function (agent, version, router) {
           if (trace) {
             spanEventRecorder = trace.traceBlockBegin()
             spanEventRecorder.recordServiceType(ServiceTypeCode.koa)
-            spanEventRecorder.recordApiDesc(`koa.router.${ctx.method.toLocaleLowerCase()} [${name}]`)
+            recordAPI(methodDescriptor, path, spanEventRecorder, ctx, name)
           }
           result = await fn.apply(this, arguments)
         } catch (e) {
           if (!e._pinpointCheck) {
             e._pinpointCheck = true
             spanEventRecorder.recordServiceType(ServiceTypeCode.koa)
-            spanEventRecorder.recordApiDesc(`koa.router.${ctx.method.toLocaleLowerCase()} [${name}]`)
+            recordAPI(methodDescriptor, path, spanEventRecorder, ctx, name)
             spanEventRecorder.recordException(e, true)
           }
           throw e
@@ -87,3 +87,14 @@ module.exports = function (agent, version, router) {
 
   return router
 }
+
+function recordAPI(methodDescriptor, path, spanEventRecorder, ctx, name) {
+  if (methodDescriptor && typeof path === 'string') {
+    spanEventRecorder.recordApiWithParameters(methodDescriptor, [path])
+  } else if (methodDescriptor) {
+    spanEventRecorder.recordApi(methodDescriptor)
+  } else {
+    spanEventRecorder.recordApiDesc(`koa.router.${ctx.method.toLocaleLowerCase()} [${name}]`)
+  }
+}
+
diff --git a/test/instrumentation/module/koa.test.js b/test/instrumentation/module/koa.test.js
index d191b348..d4a06bf5 100644
--- a/test/instrumentation/module/koa.test.js
+++ b/test/instrumentation/module/koa.test.js
@@ -25,9 +25,6 @@ const testName1 = 'koa-router1'
 test(`${testName1} Should record request in basic route koa.test.js`, function (t) {
   agent.bindHttp()
   const testName = testName1
-
-  t.plan(3)
-
   const PATH = `/${testName}`
   const app = new Koa()
   const router = new Router()
@@ -38,10 +35,11 @@ test(`${testName1} Should record request in basic route koa.test.js`, function (
     agent.callbackTraceClose((trace) => {
       let actualBuilder = new MethodDescriptorBuilder('koa', 'get')
         .setParameterDescriptor('(ctx, next)')
-        .setLineNumber(35)
+        .setLineNumber(32)
         .setFileName('koa.test.js')
       const actualMethodDescriptor = apiMetaService.cacheApiWithBuilder(actualBuilder)
       let spanEvent = trace.storage.storage[0]
+      t.equal(actualMethodDescriptor.apiId, spanEvent.apiId, 'apiId')
     })
   })
   router.post(PATH, async (ctx, next) => {
@@ -61,6 +59,7 @@ test(`${testName1} Should record request in basic route koa.test.js`, function (
     log.debug(traceMap.size)
     t.ok(traceMap.size > 0)
 
+    t.end()
     server.close()
   })
 })