More fine-tuning

starlette/middleware/http.py

@@ -7,8 +7,22 @@
 from ..types import ASGIApp, Message, Receive, Scope, Send
 
 _DispatchFlow = Union[
+ # Default case:
+ # response = yield
  AsyncGenerator[None, Response],
- AsyncGenerator[Response, Response],
+ # A/ Early response:
+ # if condition:
+ # yield Response(...)
+ # return
+ # response = yield None
+ #
+ # or B/ Error handling:
+ # try:
+ # response = yield None
+ # except ...:
+ # yield Response(...)
+ # else:
+ # ...
  AsyncGenerator[Optional[Response], Response],
 ]
 

tests/middleware/test_http.py

@@ -154,28 +154,41 @@ def test_middleware_repr():
 
 
 def test_early_response(test_client_factory):
+ async def index(request):
+ return PlainTextResponse("Hello, world!")
+
  class CustomMiddleware(HTTPMiddleware):
  async def dispatch(
  self, conn: HTTPConnection
- ) -> AsyncGenerator[Response, Response]:
- yield Response(status_code=401)
+ ) -> AsyncGenerator[Optional[Response], Response]:
+ if conn.headers.get("X-Early") == "true":
+ yield Response(status_code=401)
+ return
 
- app = Starlette(middleware=[Middleware(CustomMiddleware)])
+ yield None
+
+ app = Starlette(
+ routes=[Route("/", index)],
+ middleware=[Middleware(CustomMiddleware)],
+ )
 
  client = test_client_factory(app)
  response = client.get("/")
+ assert response.status_code == 200
+ assert response.text == "Hello, world!"
+ response = client.get("/", headers={"X-Early": "true"})
  assert response.status_code == 401
 
 
 def test_too_many_yields(test_client_factory) -> None:
- class BadMiddleware(HTTPMiddleware):
+ class CustomMiddleware(HTTPMiddleware):
  async def dispatch(
  self, conn: HTTPConnection
  ) -> AsyncGenerator[None, Response]:
  _ = yield
  yield
 
- app = Starlette(middleware=[Middleware(BadMiddleware)])
+ app = Starlette(middleware=[Middleware(CustomMiddleware)])
 
  client = test_client_factory(app)
  with pytest.raises(RuntimeError, match="should yield exactly once"):
@@ -189,7 +202,7 @@ class Failed(Exception):
  async def failure(request):
  raise Failed()
 
- class ErrorMiddleware(HTTPMiddleware):
+ class CustomMiddleware(HTTPMiddleware):
  async def dispatch(
  self, conn: HTTPConnection
  ) -> AsyncGenerator[Optional[Response], Response]:
@@ -200,7 +213,7 @@ async def dispatch(
 
  app = Starlette(
  routes=[Route("/fail", failure)],
- middleware=[Middleware(ErrorMiddleware)],
+ middleware=[Middleware(CustomMiddleware)],
  )
 
  client = test_client_factory(app)
@@ -216,7 +229,7 @@ class Failed(Exception):
  async def index(request):
  raise Failed()
 
- class BadMiddleware(HTTPMiddleware):
+ class CustomMiddleware(HTTPMiddleware):
  async def dispatch(
  self, conn: HTTPConnection
  ) -> AsyncGenerator[None, Response]:
@@ -225,7 +238,10 @@ async def dispatch(
  except Failed:
  pass
 
- app = Starlette(routes=[Route("/", index)], middleware=[Middleware(BadMiddleware)])
+ app = Starlette(
+ routes=[Route("/", index)],
+ middleware=[Middleware(CustomMiddleware)],
+ )
 
  client = test_client_factory(app)
  with pytest.raises(RuntimeError, match="no response was returned"):
@@ -260,9 +276,6 @@ async def dispatch(self, conn: HTTPConnection) -> AsyncGenerator[None, Response]
  ],
 )
 def test_contextvars(test_client_factory, middleware_cls: type):
- # this has to be an async endpoint because Starlette calls run_in_threadpool
- # on sync endpoints which has it's own set of peculiarities w.r.t propagating
- # contextvars (it propagates them forwards but not backwards)
  async def homepage(request):
  assert ctxvar.get() == "set by middleware"
  ctxvar.set("set by endpoint")