From d587eab8808fcd8b709b3186fcc2d564f5d56f2e Mon Sep 17 00:00:00 2001
From: Hiromu OCHIAI <otiai10@gmail.com>
Date: Mon, 4 Sep 2023 00:31:16 +0900
Subject: [PATCH] Refactor onerror handler

---
 src/router.ts             | 21 +++++++++++++++++----
 tests/spec/router.spec.ts | 18 ++++++++++++++++++
 2 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/src/router.ts b/src/router.ts
index 2dce611..f2033bc 100644
--- a/src/router.ts
+++ b/src/router.ts
@@ -33,6 +33,16 @@ export class Router<T extends chrome.events.Event<any>, U = Record<string, unkno
     this.notfound = callback
   }
 
+  // Error handler
+  private error: HandlerOf<ExtractCallback<T>> = async function (this: { route: Resolved<U> }, ...args) {
+    return { status: 500, message: `Handler for request "${this.route[ActionKey]}" throw an error` }
+  }
+
+  // onError can overwire behavior for 500-ish error
+  public onError (callback: HandlerOf<ExtractCallback<T>>): void {
+    this.error = callback
+  }
+
   private readonly routes: {
     exact: Array<RouteMatcher<HandlerOf<ExtractCallback<T>>>>
     regex: Array<RouteMatcher<HandlerOf<ExtractCallback<T>>>>
@@ -80,11 +90,14 @@ export class Router<T extends chrome.events.Event<any>, U = Record<string, unkno
       this.resolver(...args).then(route => {
         const fn = this.findHandler(route[ActionKey])
         const res = fn(...args)
-        if (res instanceof Promise) void res.then(sendResponse)
-        else sendResponse(res)
+        if (res instanceof Promise) {
+          res.then(sendResponse).catch(err => this.error.bind({
+            route, error: err
+          })(...args).then(sendResponse))
+        } else sendResponse(res)
       }).catch(err => {
-        console.error(err)
-        sendResponse({ [ActionKey]: '__notfound__', ...args }) // TODO: Handle error case
+        const fn = this.error.bind({ route: { [ActionKey]: '__router_error__', error: err } })
+        void fn(...args).then(sendResponse)
       })
       return true
     }) as ExtractCallback<T>
diff --git a/tests/spec/router.spec.ts b/tests/spec/router.spec.ts
index 13a389d..a53432d 100644
--- a/tests/spec/router.spec.ts
+++ b/tests/spec/router.spec.ts
@@ -70,4 +70,22 @@ describe('Router', () => {
       expect(sendResponse).toBeCalledWith({ status: 5004, message: 'See you yesterday ;)' })
     })
   })
+
+  describe('onError', () => {
+    it('should overwrite a handler for error', async () => {
+      const r = new Router<chrome.runtime.ExtensionMessageEvent, { status: number }>()
+      const onerror = jest.fn().mockName('callback').mockImplementation(async function (this: any) {
+        return { message: 'See you yesterday ;)', status: 500, error: this.error }
+      })
+      const sendResponse = jest.fn().mockName('sendResponse')
+      r.onError(onerror)
+      r.on('/problem', async function () {
+        throw new Error('Something wrong')
+      })
+      r.listener()({ action: '/problem' }, {}, sendResponse)
+      await new Promise(resolve => setTimeout(resolve, 0))
+      expect(onerror).toBeCalled()
+      expect(sendResponse).toBeCalledWith({ status: 500, message: 'See you yesterday ;)', error: new Error('Something wrong') })
+    })
+  })
 })