add typehint and tests for HandlerChain.respond

rolo/gateway/chain.py

@@ -188,7 +188,12 @@ def handle(self, context: RC, response: Response):
             if not self.finalized:
                 self._call_finalizers(response)
 
-    def respond(self, status_code: int = 200, payload: t.Any = None, headers: Headers = None):
+    def respond(
+        self,
+        status_code: int = 200,
+        payload: t.Any = None,
+        headers: t.Union[Headers, t.Mapping] = None,
+    ):
         """
         Convenience method for handlers to stop the chain and set the given status and payload to the
         current response object.

tests/gateway/test_chain.py

@@ -1,5 +1,7 @@
 from unittest import mock
 
+from werkzeug.datastructures import Headers
+
 from rolo.gateway import CompositeFinalizer, CompositeHandler, HandlerChain, RequestContext
 from rolo.response import Response
 
@@ -84,6 +86,32 @@ def _raise(*args, **kwargs):
     assert chain.error is None
 
 
+def test_respond_with_json_response():
+    def handle(chain_: HandlerChain, _context, _response):
+        chain_.respond(202, {"foo": "bar"}, headers={"X-Foo": "Bar"})
+
+    chain = HandlerChain(request_handlers=[handle])
+    chain.handle(RequestContext(), Response())
+
+    assert chain.response.status_code == 202
+    assert chain.response.json == {"foo": "bar"}
+    assert chain.response.headers.get("x-foo") == "Bar"
+    assert chain.response.mimetype == "application/json"
+
+
+def test_respond_with_string_response():
+    def handle(chain_: HandlerChain, _context, _response):
+        chain_.respond(200, "foobar", Headers({"X-Foo": "Bar"}))
+
+    chain = HandlerChain(request_handlers=[handle])
+    chain.handle(RequestContext(), Response())
+
+    assert chain.response.status_code == 200
+    assert chain.response.data == b"foobar"
+    assert chain.response.headers.get("x-foo") == "Bar"
+    assert chain.response.mimetype == "text/plain"
+
+
 class TestCompositeHandler:
     def test_composite_handler_stops_handler_chain(self):
         def inner1(_chain: HandlerChain, request: RequestContext, response: Response):