chore: add async middleware tests

open-telemetry · Jul 14, 2020 · a9af954 · a9af954
1 parent 206dd8f
commit a9af954
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 3 deletions.
diff --git a/plugins/node/opentelemetry-plugin-koa/src/koa.ts b/plugins/node/opentelemetry-plugin-koa/src/koa.ts
@@ -116,8 +116,7 @@ export class KoaPlugin extends BasePlugin<typeof koa> {
     if (middlewareLayer[kLayerPatched] === true) return middlewareLayer;
     middlewareLayer[kLayerPatched] = true;
     this._logger.debug('patching Koa middleware layer');
-
-    return (context: KoaContext, next: koa.Next) => {
+    return async (context: KoaContext, next: koa.Next) => {
       if (this._tracer.getCurrentSpan() === undefined) {
         return middlewareLayer(context, next);
       }
@@ -132,7 +131,7 @@ export class KoaPlugin extends BasePlugin<typeof koa> {
       });
       const startTime = hrTime();
 
-      const result = middlewareLayer(context, next);
+      const result = await middlewareLayer(context, next);
       span.end(startTime);
       return result;
     };

diff --git a/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts b/plugins/node/opentelemetry-plugin-koa/test/koa.test.ts
@@ -83,6 +83,13 @@ describe('Koa Plugin - Core Tests', () => {
     return next();
   };
 
+  const asyncMiddleware: koa.Middleware = async (ctx, next) => {
+    const start = Date.now();
+    await next();
+    const ms = Date.now() - start;
+    console.log(`${ctx.method} ${ctx.url} - ${ms}ms`);
+  };
+
   describe('Instrumenting core middleware calls', () => {
     it('should create a child span for middlewares', async () => {
       const rootSpan = tracer.startSpan('rootSpan');
@@ -151,6 +158,41 @@ describe('Koa Plugin - Core Tests', () => {
       assert.strictEqual(res, 'test');
       server.close();
     });
+
+    it('should handle async middleware functions', async () => {
+      const rootSpan = tracer.startSpan('rootSpan');
+      const app = new koa();
+      app.use((ctx, next) => tracer.withSpan(rootSpan, next));
+      app.use(asyncMiddleware);
+
+      const server = http.createServer(app.callback());
+      await new Promise(resolve => server.listen(0, resolve));
+      const port = (server.address() as AddressInfo).port;
+      assert.strictEqual(memoryExporter.getFinishedSpans().length, 0);
+
+      await tracer.withSpan(rootSpan, async () => {
+        await httpRequest.get(`http://localhost:${port}`);
+        rootSpan.end();
+
+        const requestHandlerSpan = memoryExporter
+          .getFinishedSpans()
+          .find(span => span.name.includes('asyncMiddleware'));
+        assert.notStrictEqual(requestHandlerSpan, undefined);
+        assert.strictEqual(
+          requestHandlerSpan?.attributes[AttributeNames.COMPONENT],
+          KoaComponentName
+        );
+        assert.strictEqual(
+          requestHandlerSpan?.attributes[AttributeNames.KOA_TYPE],
+          KoaLayerType.MIDDLEWARE
+        );
+        const exportedRootSpan = memoryExporter
+          .getFinishedSpans()
+          .find(span => span.name === 'rootSpan');
+        assert.notStrictEqual(exportedRootSpan, undefined);
+      });
+      server.close();
+    });
   });
 
   describe('Disabling plugin', () => {