diff --git a/test/instrumentation/module/koa.test.js b/test/instrumentation/module/koa.test.js
index b1f98d8c..d191b348 100644
--- a/test/instrumentation/module/koa.test.js
+++ b/test/instrumentation/module/koa.test.js
@@ -35,13 +35,13 @@ test(`${testName1} Should record request in basic route koa.test.js`, function (
   router.get(PATH, async (ctx, next) => {
     ctx.body = 'ok. get'
 
-    const trace = agent.traceContext.currentTraceObject()
-    setImmediate(() => {
+    agent.callbackTraceClose((trace) => {
       let actualBuilder = new MethodDescriptorBuilder('koa', 'get')
         .setParameterDescriptor('(ctx, next)')
-        .setLineNumber(34)
+        .setLineNumber(35)
         .setFileName('koa.test.js')
       const actualMethodDescriptor = apiMetaService.cacheApiWithBuilder(actualBuilder)
+      let spanEvent = trace.storage.storage[0]
     })
   })
   router.post(PATH, async (ctx, next) => {
diff --git a/test/support/agent-singleton-mock.js b/test/support/agent-singleton-mock.js
index 6a851904..8ac613da 100644
--- a/test/support/agent-singleton-mock.js
+++ b/test/support/agent-singleton-mock.js
@@ -77,9 +77,15 @@ class MockAgent extends Agent {
         shimmer.unwrap(http, 'request')
     }
 
-    resetAgent(callback) {
-        this.pinpointClient = new MockPinpointClient(this.config, this.agentInfo, this.dataSender)
+    callbackTraceClose(callback) {
+        const trace = this.traceContext.currentTraceObject()
+        const origin = trace.close
+        trace.close = () => {
+            callback(trace)
+            origin.apply(trace, arguments)
+        }
     }
+
 }
 
 const agent = new MockAgent(fixture.config)