From ed5ca5e756ce8f8dc7f94525c2cc9d73093fa0b4 Mon Sep 17 00:00:00 2001
From: Jonatan Nilsson <jonatan@nilsson.is>
Date: Mon, 11 Oct 2021 02:06:22 +0000
Subject: [PATCH] hopefully fix so after gets called after every request

---
 flaska.mjs               |  6 +++++-
 test.mjs                 | 10 ++++++++--
 test/flaska.api.test.mjs |  1 -
 test/flaska.in.test.mjs  | 15 +++++++++++----
 4 files changed, 24 insertions(+), 8 deletions(-)

diff --git a/flaska.mjs b/flaska.mjs
index eeb3b25..7a43ff2 100644
--- a/flaska.mjs
+++ b/flaska.mjs
@@ -276,7 +276,6 @@ export class Flaska {
       ctx.log.error(err)
       ctx.res.statusCode = ctx.statusCode = 400
       ctx.res.end()
-      ctx.finished = true
     }
     this._onreserror = function(err, ctx) {
       ctx.log.error(err)
@@ -405,6 +404,9 @@ export class Flaska {
     req.on('close', () => {
       this.requestEnded(ctx)
     })
+    res.on('finish', () => {
+      this.requestEnded(ctx)
+    })
 
     try {
       this._beforeCompiled(ctx)
@@ -496,6 +498,8 @@ export class Flaska {
   }
 
   requestEnded(ctx) {
+    if (ctx.finished) return
+    ctx.finished = true
     this._afterCompiled(ctx)
     if (this._afterAsyncCompiled) {
       return this._afterAsyncCompiled(ctx).then()
diff --git a/test.mjs b/test.mjs
index 5d76f28..4758d71 100644
--- a/test.mjs
+++ b/test.mjs
@@ -7,11 +7,17 @@ flaska.get('/', function(ctx) {
   return new Promise(function(res, rej) {
     ctx.body = { status: true }
     setTimeout(res, 1000)
-  }).then(function() {
-    process.stdout.write(".")
   })
 })
 
+flaska.after(function(ctx) {
+  if (ctx.aborted) {
+    process.stdout.write("A")
+  } else {
+    process.stdout.write(".")
+  }
+})
+
 flaska.listen(port, function() {
   console.log('listening on port', port)
 })
diff --git a/test/flaska.api.test.mjs b/test/flaska.api.test.mjs
index f7d2bda..9da12fb 100644
--- a/test/flaska.api.test.mjs
+++ b/test/flaska.api.test.mjs
@@ -84,7 +84,6 @@ specialHandlers.forEach(function(type) {
         let ctx = createCtx()
         flaska['_' + type](new Error(), ctx)
         assert.strictEqual(ctx.res.statusCode, 400)
-        assert.strictEqual(ctx.finished, true)
         assert.strictEqual(ctx.res.end.callCount, 1)
         assert.strictEqual(ctx.res.end.firstCall.length, 0)
       })
diff --git a/test/flaska.in.test.mjs b/test/flaska.in.test.mjs
index 73534b1..7584956 100644
--- a/test/flaska.in.test.mjs
+++ b/test/flaska.in.test.mjs
@@ -23,7 +23,7 @@ t.describe('#requestStart()', function() {
       try {
         assert.ok(err)
         assert.strictEqual(assertReq.on.callCount, 3)
-        assert.strictEqual(assertRes.on.callCount, 1)
+        assert.strictEqual(assertRes.on.callCount, 2)
 
 
         assert.strictEqual(assertRes.on.firstCall[0], 'error')
@@ -32,6 +32,13 @@ t.describe('#requestStart()', function() {
         assert.strictEqual(onResError.callCount, 1)
         assert.strictEqual(onResError.firstCall[0], assertErrorTwo)
         assert.strictEqual(onResError.firstCall[1], ctx)
+        
+        assert.strictEqual(assertRes.on.secondCall[0], 'finish')
+        assert.strictEqual(typeof(assertRes.on.secondCall[1]), 'function')
+        assert.strictEqual(onEnded.callCount, 0)
+        assertRes.on.secondCall[1]()
+        assert.strictEqual(onEnded.callCount, 1)
+        assert.strictEqual(onEnded.firstCall[0], ctx)
 
         assert.strictEqual(assertReq.on.firstCall[0], 'error')
         assert.strictEqual(typeof(assertReq.on.firstCall[1]), 'function')
@@ -48,10 +55,10 @@ t.describe('#requestStart()', function() {
         
         assert.strictEqual(assertReq.on.thirdCall[0], 'close')
         assert.strictEqual(typeof(assertReq.on.thirdCall[1]), 'function')
-        assert.strictEqual(onEnded.called, false)
+        assert.strictEqual(onEnded.callCount, 1)
         assertReq.on.thirdCall[1]()
-        assert.strictEqual(onEnded.called, true)
-        assert.strictEqual(onEnded.firstCall[0], ctx)
+        assert.strictEqual(onEnded.callCount, 2)
+        assert.strictEqual(onEnded.secondCall[0], ctx)
 
         // Test abort and close