From b9a1e71610b2d24b8f982718cee736bf3dc58390 Mon Sep 17 00:00:00 2001
From: Jon Friesen <jon@qpoint.io>
Date: Wed, 10 Jan 2024 18:50:22 -0700
Subject: [PATCH] awaits on complete fn (#7)

---
 src/endpoint-harness.ts | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/src/endpoint-harness.ts b/src/endpoint-harness.ts
index 1567aca..63b80f6 100644
--- a/src/endpoint-harness.ts
+++ b/src/endpoint-harness.ts
@@ -47,12 +47,6 @@ export class EndpointHarness {
       passThroughOnException: () => { }
     };
     this.reqCtx = new Context(this.req, {}, this.ctx, {});
-
-    // extract the context reference
-    this.endpoint.use((ctx: Context, next: Function) => {
-      this.reqCtx = ctx;
-      next();
-    });
   }
 
   use(fn: Function): this {
@@ -66,12 +60,18 @@ export class EndpointHarness {
     const defaultEnv = env || {};
     const defaultCtx = ctx || this.ctx
 
+    // extract the final context reference
+    this.endpoint.use(async (ctx: Context, next: Function) => {
+      this.reqCtx = ctx;
+      await next();
+    });
+
     return this.endpoint.fetch(defaultRequest, defaultEnv, defaultCtx);
   }
 
   async onAfterComplete(fn: Function) {
     await this.queue.allComplete();
-    fn(this.reqCtx);
+    await fn(this.reqCtx);
   }
 }