Fix #193 by enabling listener middleware to return a response

seratch · seratch · commit 55d4f16975fd · 2021-01-07T12:14:28.000+09:00
diff --git a/slack_bolt/app/app.py b/slack_bolt/app/app.py
@@ -2,6 +2,7 @@
 import json
 import logging
 import os
+import time
 from concurrent.futures.thread import ThreadPoolExecutor
 from http.server import SimpleHTTPRequestHandler, HTTPServer
 from typing import List, Union, Pattern, Callable, Dict, Optional, Sequence
@@ -42,6 +43,7 @@
     error_client_invalid_type,
     error_authorize_conflicts,
     warning_bot_only_conflicts,
+    debug_return_listener_middleware_response,
 )
 from slack_bolt.middleware import (
     Middleware,
@@ -323,6 +325,7 @@ def dispatch(self, req: BoltRequest) -> BoltResponse:
         :param req: An incoming request from Slack.
         :return: The response generated by this Bolt app.
         """
+        starting_time = time.time()
         self._init_context(req)
 
         resp: BoltResponse = BoltResponse(status=200, body="")
@@ -348,12 +351,27 @@ def middleware_next():
             self._framework_logger.debug(debug_checking_listener(listener_name))
             if listener.matches(req=req, resp=resp):
                 # run all the middleware attached to this listener first
-                resp, next_was_not_called = listener.run_middleware(req=req, resp=resp)
+                middleware_resp, next_was_not_called = listener.run_middleware(
+                    req=req, resp=resp
+                )
                 if next_was_not_called:
+                    if middleware_resp is not None:
+                        if self._framework_logger.level <= logging.DEBUG:
+                            debug_message = debug_return_listener_middleware_response(
+                                listener_name,
+                                middleware_resp.status,
+                                middleware_resp.body,
+                                starting_time,
+                            )
+                            self._framework_logger.debug(debug_message)
+                        return middleware_resp
                     # The last listener middleware didn't call next() method.
                     # This means the listener is not for this incoming request.
                     continue
 
+                if middleware_resp is not None:
+                    resp = middleware_resp
+
                 self._framework_logger.debug(debug_running_listener(listener_name))
                 listener_response: Optional[BoltResponse] = self._listener_runner.run(
                     request=req,
diff --git a/slack_bolt/app/async_app.py b/slack_bolt/app/async_app.py
@@ -1,6 +1,7 @@
 import inspect
 import logging
 import os
+import time
 from typing import Optional, List, Union, Callable, Pattern, Dict, Awaitable, Sequence
 
 from aiohttp import web
@@ -39,6 +40,7 @@
     error_oauth_settings_invalid_type_async,
     error_oauth_flow_invalid_type_async,
     warning_bot_only_conflicts,
+    debug_return_listener_middleware_response,
 )
 from slack_bolt.lazy_listener.asyncio_runner import AsyncioLazyListenerRunner
 from slack_bolt.listener.async_listener import AsyncListener, AsyncCustomListener
@@ -359,6 +361,7 @@ async def async_dispatch(self, req: AsyncBoltRequest) -> BoltResponse:
         :param req: An incoming request from Slack.
         :return: The response generated by this Bolt app.
         """
+        starting_time = time.time()
         self._init_context(req)
 
         resp: BoltResponse = BoltResponse(status=200, body="")
@@ -386,14 +389,28 @@ async def async_middleware_next():
             self._framework_logger.debug(debug_checking_listener(listener_name))
             if await listener.async_matches(req=req, resp=resp):
                 # run all the middleware attached to this listener first
-                resp, next_was_not_called = await listener.run_async_middleware(
-                    req=req, resp=resp
-                )
+                (
+                    middleware_resp,
+                    next_was_not_called,
+                ) = await listener.run_async_middleware(req=req, resp=resp)
                 if next_was_not_called:
+                    if middleware_resp is not None:
+                        if self._framework_logger.level <= logging.DEBUG:
+                            debug_message = debug_return_listener_middleware_response(
+                                listener_name,
+                                middleware_resp.status,
+                                middleware_resp.body,
+                                starting_time,
+                            )
+                            self._framework_logger.debug(debug_message)
+                        return middleware_resp
                     # The last listener middleware didn't call next() method.
                     # This means the listener is not for this incoming request.
                     continue
 
+                if middleware_resp is not None:
+                    resp = middleware_resp
+
                 self._framework_logger.debug(debug_running_listener(listener_name))
                 listener_response: Optional[
                     BoltResponse
diff --git a/slack_bolt/listener/async_listener.py b/slack_bolt/listener/async_listener.py
@@ -33,7 +33,7 @@ async def run_async_middleware(
         *,
         req: AsyncBoltRequest,
         resp: BoltResponse,
-    ) -> Tuple[BoltResponse, bool]:
+    ) -> Tuple[Optional[BoltResponse], bool]:
         """Runs an async middleware.
 
         :param req: The incoming request
diff --git a/slack_bolt/listener/asyncio_runner.py b/slack_bolt/listener/asyncio_runner.py
@@ -42,9 +42,10 @@ async def run(
         response: BoltResponse,
         listener_name: str,
         listener: AsyncListener,
+        starting_time: Optional[float] = None,
     ) -> Optional[BoltResponse]:
         ack = request.context.ack
-        starting_time = time.time()
+        starting_time = starting_time if starting_time is not None else time.time()
         if self.process_before_response:
             if not request.lazy_only:
                 try:
diff --git a/slack_bolt/listener/listener.py b/slack_bolt/listener/listener.py
@@ -1,5 +1,5 @@
 from abc import abstractmethod, ABCMeta
-from typing import Callable, Tuple, Sequence
+from typing import Callable, Tuple, Sequence, Optional
 
 from slack_bolt.listener_matcher import ListenerMatcher
 from slack_bolt.middleware import Middleware
@@ -32,7 +32,7 @@ def run_middleware(
         *,
         req: BoltRequest,
         resp: BoltResponse,
-    ) -> Tuple[BoltResponse, bool]:
+    ) -> Tuple[Optional[BoltResponse], bool]:
         """
 
         :param req: the incoming request
diff --git a/slack_bolt/listener/thread_runner.py b/slack_bolt/listener/thread_runner.py
@@ -43,9 +43,10 @@ def run(  # type: ignore
         response: BoltResponse,
         listener_name: str,
         listener: Listener,
+        starting_time: Optional[float] = None,
     ) -> Optional[BoltResponse]:
         ack = request.context.ack
-        starting_time = time.time()
+        starting_time = starting_time if starting_time is not None else time.time()
         if self.process_before_response:
             if not request.lazy_only:
                 try:
diff --git a/slack_bolt/logger/messages.py b/slack_bolt/logger/messages.py
@@ -1,3 +1,4 @@
+import time
 from typing import Union
 
 from slack_sdk.web import SlackResponse
@@ -111,3 +112,10 @@ def debug_running_lazy_listener(func_name: str) -> str:
 
 def debug_responding(status: int, body: str, millis: int) -> str:
     return f'Responding with status: {status} body: "{body}" ({millis} millis)'
+
+
+def debug_return_listener_middleware_response(
+    listener_name: str, status: int, body: str, starting_time: float
+) -> str:
+    millis = int((time.time() - starting_time) * 1000)
+    return f"Responding with listener middleware's response - listener: {listener_name}, status: {status}, body: {body} ({millis} millis)"
diff --git a/slack_bolt/middleware/async_middleware.py b/slack_bolt/middleware/async_middleware.py
@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import Callable, Awaitable
+from typing import Callable, Awaitable, Optional
 
 from slack_bolt.request.async_request import AsyncBoltRequest
 from slack_bolt.response import BoltResponse
@@ -13,7 +13,7 @@ async def async_process(
         req: AsyncBoltRequest,
         resp: BoltResponse,
         next: Callable[[], Awaitable[BoltResponse]],
-    ) -> BoltResponse:
+    ) -> Optional[BoltResponse]:
         raise NotImplementedError()
 
     @property
diff --git a/slack_bolt/middleware/middleware.py b/slack_bolt/middleware/middleware.py
@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import Callable
+from typing import Callable, Optional
 
 from slack_bolt.request import BoltRequest
 from slack_bolt.response import BoltResponse
@@ -13,7 +13,7 @@ def process(
         req: BoltRequest,
         resp: BoltResponse,
         next: Callable[[], BoltResponse],
-    ) -> BoltResponse:
+    ) -> Optional[BoltResponse]:
         raise NotImplementedError()
 
     @property
diff --git a/slack_bolt/workflows/step/step_middleware.py b/slack_bolt/workflows/step/step_middleware.py
@@ -19,7 +19,7 @@ def process(
         req: BoltRequest,
         resp: BoltResponse,
         next: Callable[[], BoltResponse],
-    ) -> BoltResponse:
+    ) -> Optional[BoltResponse]:
 
         if self.step.edit.matches(req=req, resp=resp):
             resp = self._run(self.step.edit, req, resp)
diff --git a/tests/scenario_tests/test_listener_middleware.py b/tests/scenario_tests/test_listener_middleware.py
@@ -0,0 +1,102 @@
+import json
+from time import time
+
+from slack_sdk.signature import SignatureVerifier
+from slack_sdk.web import WebClient
+
+from slack_bolt import BoltResponse
+from slack_bolt.app import App
+from slack_bolt.request import BoltRequest
+from tests.mock_web_api_server import (
+    setup_mock_web_api_server,
+    cleanup_mock_web_api_server,
+)
+from tests.utils import remove_os_env_temporarily, restore_os_env
+
+
+class TestListenerMiddleware:
+    signing_secret = "secret"
+    valid_token = "xoxb-valid"
+    mock_api_server_base_url = "http://localhost:8888"
+    signature_verifier = SignatureVerifier(signing_secret)
+    web_client = WebClient(
+        token=valid_token,
+        base_url=mock_api_server_base_url,
+    )
+
+    def setup_method(self):
+        self.old_os_env = remove_os_env_temporarily()
+        setup_mock_web_api_server(self)
+
+    def teardown_method(self):
+        cleanup_mock_web_api_server(self)
+        restore_os_env(self.old_os_env)
+
+    body = {
+        "type": "shortcut",
+        "token": "verification_token",
+        "action_ts": "111.111",
+        "team": {
+            "id": "T111",
+            "domain": "workspace-domain",
+            "enterprise_id": "E111",
+            "enterprise_name": "Org Name",
+        },
+        "user": {"id": "W111", "username": "primary-owner", "team_id": "T111"},
+        "callback_id": "test-shortcut",
+        "trigger_id": "111.111.xxxxxx",
+    }
+
+    def build_request(self) -> BoltRequest:
+        timestamp, body = str(int(time())), json.dumps(self.body)
+        return BoltRequest(
+            body=body,
+            headers={
+                "content-type": ["application/json"],
+                "x-slack-signature": [
+                    self.signature_verifier.generate_signature(
+                        body=body,
+                        timestamp=timestamp,
+                    )
+                ],
+                "x-slack-request-timestamp": [timestamp],
+            },
+        )
+
+    def test_return_response(self):
+        app = App(
+            client=self.web_client,
+            signing_secret=self.signing_secret,
+        )
+
+        @app.shortcut(
+            constraints="test-shortcut",
+            middleware=[listener_middleware_returning_response],
+        )
+        def handle(ack):
+            ack()
+
+        response = app.dispatch(self.build_request())
+        assert response.status == 200
+        assert response.body == "listener middleware"
+
+    def test_next(self):
+        app = App(
+            client=self.web_client,
+            signing_secret=self.signing_secret,
+        )
+
+        @app.shortcut(constraints="test-shortcut", middleware=[just_next])
+        def handle(ack):
+            ack()
+
+        response = app.dispatch(self.build_request())
+        assert response.status == 200
+
+
+def listener_middleware_returning_response():
+    return BoltResponse(status=200, body="listener middleware")
+
+
+def just_next(next):
+    next()
diff --git a/tests/scenario_tests_async/test_listener_middleware.py b/tests/scenario_tests_async/test_listener_middleware.py
