Simple test case to check if header is available in failed request response, if handled

tests/test_context/test_context_on_unhandled_exception.py

@@ -5,15 +5,26 @@
 from starlette.responses import Response
 from starlette.testclient import TestClient
 
-from starlette_context import middleware, plugins
+from starlette_context import middleware, plugins, context
 from starlette_context.header_keys import HeaderKeys
 
+
+class CustomExc(Exception):
+    pass
+
+
+def custom_exc_handler(request: Request, exc):
+    return Response(status_code=400, headers=context.data)
+
+
 app = Starlette(
     middleware=[
         Middleware(
-            middleware.ContextMiddleware, plugins=(plugins.RequestIdPlugin(),)
+            middleware.RawContextMiddleware,
+            plugins=(plugins.RequestIdPlugin(),),
         )
-    ]
+    ],
+    exception_handlers={CustomExc: custom_exc_handler},
 )
 
 
@@ -22,6 +33,11 @@ async def exc(request: Request):
     raise Exception
 
 
+@app.route("/handled")
+async def handled(request: Request):
+    raise CustomExc
+
+
 @app.route("/")
 async def index(request: Request):
     return Response(status_code=status.HTTP_200_OK)
@@ -38,3 +54,7 @@ def test_context_persistence():
     resp = client.get("/exc")
     assert resp.status_code == 500
     assert HeaderKeys.request_id.lower() not in resp.headers
+
+    resp = client.get("/handled")
+    assert resp.status_code == 400
+    assert HeaderKeys.request_id.lower() in resp.headers