Minor feature requests (#13)

* add test for headers in response when exception is raised in view

* Check if context already filled (#12)

* add tests for checking if context is available

* extract base for uuid plugins, add force_new_uuid, add uuid validation

Co-authored-by: Viacheslav Kukushkin <vy.kukushkin@gmail.com>

Author: tomwojcik

starlette_context/ctx.py

@@ -1,6 +1,7 @@
 from collections import UserDict
 from typing import Any
 
+from contextvars import copy_context
 from starlette_context import _request_scope_context_storage
 
 
@@ -32,6 +33,9 @@ def data(self) -> dict:
                 "outside of the request-response cycle."
             ) from e
 
+    def exists(self) -> bool:
+        return _request_scope_context_storage in copy_context()
+
     def copy(self) -> dict:
         """
         Read only context data.

starlette_context/plugins/correlation_id.py

@@ -1,25 +1,6 @@
-import uuid
-from typing import Optional
-
-from starlette.requests import Request
-from starlette.responses import Response
-
 from starlette_context.header_keys import HeaderKeys
-from starlette_context.plugins.plugin import Plugin
+from starlette_context.plugins.plugin_uuid import PluginUUIDBase
 
 
-class CorrelationIdPlugin(Plugin):
+class CorrelationIdPlugin(PluginUUIDBase):
     key = HeaderKeys.correlation_id
-
-    async def extract_value_from_header_by_key(
-        self, request: Request
-    ) -> Optional[str]:
-        await super(
-            CorrelationIdPlugin, self
-        ).extract_value_from_header_by_key(request)
-        if self.value is None:
-            self.value = uuid.uuid4().hex
-        return self.value
-
-    async def enrich_response(self, response: Response) -> None:
-        await self._add_kv_to_response_headers(response)

starlette_context/plugins/plugin_uuid.py

@@ -0,0 +1,44 @@
+import uuid
+from typing import Optional
+
+from starlette.requests import Request
+from starlette.responses import Response
+
+from starlette_context.plugins.plugin import Plugin
+
+
+class PluginUUIDBase(Plugin):
+    uuid_functions_mapper = {4: uuid.uuid4}
+
+    def __init__(self, force_new_uuid: bool = False, version: int = 4):
+        super().__init__()
+        if version not in self.uuid_functions_mapper:
+            raise TypeError(f"UUID version {version} is not supported.")
+        self.force_new_uuid = force_new_uuid
+        self.version = version
+
+    def validate_uuid(self, uuid_to_validate: str) -> None:
+        try:
+            uuid.UUID(uuid_to_validate, version=self.version)
+        except Exception as e:
+            raise ValueError("Wrong uuid") from e
+
+    def get_new_uuid(self) -> str:
+        func = self.uuid_functions_mapper[self.version]
+        return func().hex
+
+    async def extract_value_from_header_by_key(
+        self, request: Request
+    ) -> Optional[str]:
+        if not self.force_new_uuid:
+            await super().extract_value_from_header_by_key(request)
+
+        if self.value is None:
+            self.value = self.get_new_uuid()
+
+        self.validate_uuid(self.value)
+
+        return self.value
+
+    async def enrich_response(self, response: Response) -> None:
+        await self._add_kv_to_response_headers(response)

starlette_context/plugins/request_id.py

@@ -1,25 +1,6 @@
-import uuid
-from typing import Optional
-
-from starlette.requests import Request
-from starlette.responses import Response
-
 from starlette_context.header_keys import HeaderKeys
-from starlette_context.plugins.plugin import Plugin
+from starlette_context.plugins.plugin_uuid import PluginUUIDBase
 
 
-class RequestIdPlugin(Plugin):
+class RequestIdPlugin(PluginUUIDBase):
     key = HeaderKeys.request_id
-
-    async def extract_value_from_header_by_key(
-        self, request: Request
-    ) -> Optional[str]:
-        await super(RequestIdPlugin, self).extract_value_from_header_by_key(
-            request
-        )
-        if self.value is None:
-            self.value = uuid.uuid4().hex
-        return self.value
-
-    async def enrich_response(self, response: Response) -> None:
-        await self._add_kv_to_response_headers(response)

tests/conftest.py

@@ -7,9 +7,10 @@
 
 from starlette_context import middleware
 from starlette_context.header_keys import HeaderKeys
+import uuid
 
-dummy_correlation_id = "dummy_correlation_id"
-dummy_request_id = "dummy_request_id"
+dummy_correlation_id = uuid.uuid4().hex
+dummy_request_id = uuid.uuid4().hex
 dummy_user_agent = "dummy_user_agent"
 dummy_date = "Wed, 01 Jan 2020 04:27:12 GMT"
 dummy_forwarded_for = "203.0.113.19"

tests/test_integration/test_context_existence.py

@@ -0,0 +1,32 @@
+import json
+
+from starlette.applications import Starlette
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response
+from starlette.testclient import TestClient
+
+from starlette_context import context, middleware
+
+
+app = Starlette()
+app.add_middleware(middleware.ContextMiddleware)
+
+
+@app.route("/")
+async def index(request: Request) -> Response:
+    return JSONResponse({"exists": context.exists()})
+
+
+client = TestClient(app)
+
+
+def test_context_existence_in_request_response_cycle():
+    resp = client.get("/")
+    assert json.loads(resp.content) == {"exists": True}
+
+
+def test_context_outside_of_request_response_cycle():
+    assert context.exists() is False
+    resp = client.get("/")
+    assert context.exists() is False

tests/test_integration/test_plugins_response_headers.py

@@ -10,6 +10,7 @@
 from starlette_context import context, plugins
 from starlette_context.middleware import ContextMiddleware
 from tests.conftest import dummy_correlation_id, dummy_request_id
+from starlette_context.header_keys import HeaderKeys
 
 middleware = [
     Middleware(
@@ -46,8 +47,8 @@ def dt_serializator(o):
 def test_response_headers(headers):
     response = client.get("/", headers=headers)
     assert 2 == len(response.headers)
-    cid_header = response.headers["x-correlation-id"]
-    rid_header = response.headers["x-request-id"]
+    cid_header = response.headers[HeaderKeys.correlation_id.lower()]
+    rid_header = response.headers[HeaderKeys.request_id.lower()]
     assert dummy_correlation_id == cid_header
     assert dummy_request_id == rid_header
 

tests/test_integration/test_server_error_middleware.py

@@ -0,0 +1,66 @@
+import json
+import uuid
+from typing import NoReturn
+
+from starlette.middleware import Middleware
+
+from starlette_context import plugins
+from starlette.applications import Starlette
+
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.testclient import TestClient
+
+from starlette_context import middleware
+from starlette_context.header_keys import HeaderKeys
+from starlette_context import context
+
+
+class CustomException(Exception):
+    pass
+
+
+async def custom_exception_handler(request: Request, exc: Exception):
+    return JSONResponse({"exception": "handled"}, headers=context.data)
+
+
+async def general_exception_handler(request: Request, exc: Exception):
+    return JSONResponse({"exception": "handled"}, headers=context.data)
+
+
+middleware = [
+    Middleware(
+        middleware.ContextMiddleware, plugins=(plugins.RequestIdPlugin(),)
+    )
+]
+exception_handlers = {
+    CustomException: custom_exception_handler,
+    Exception: general_exception_handler,
+}
+
+
+app = Starlette(exception_handlers=exception_handlers, middleware=middleware,)
+
+headers = {HeaderKeys.request_id: uuid.uuid4().hex}
+
+
+@app.route("/")
+async def index(_) -> NoReturn:
+    raise RuntimeError
+
+
+@app.route("/custom-exc")
+async def index(_) -> NoReturn:
+    raise CustomException
+
+
+client = TestClient(app)
+
+
+def test_exception_handling_that_is_not_resulting_in_500():
+    resp = client.get("/custom-exc", headers=headers)
+    assert json.loads(resp.content) == {"exception": "handled"}
+    assert (
+        resp.headers.get(HeaderKeys.request_id)
+        == headers[HeaderKeys.request_id]
+    )

tests/test_unit/test_plugins/test_correlation_id.py

@@ -20,6 +20,15 @@ async def test_process_request_for_existing_header(
     assert dummy_correlation_id == plugin.value
 
 
+@pytest.mark.asyncio
+async def test_invalid_correlation_id_uuid(
+    plugin: plugins.CorrelationIdPlugin, mocked_request: Request
+):
+    mocked_request.headers[HeaderKeys.correlation_id] = "invalid_uuid"
+    with pytest.raises(ValueError):
+        await plugin.process_request(mocked_request)
+
+
 @pytest.mark.asyncio
 async def test_process_request_for_missing_header(
     plugin: plugins.CorrelationIdPlugin, mocked_request: Request
@@ -52,3 +61,35 @@ async def test_enrich_response_str(
         dummy_correlation_id
         == mocked_response.headers[HeaderKeys.correlation_id]
     )
+
+
+def test_version_cant_map_to_function():
+    with pytest.raises(TypeError):
+        plugins.CorrelationIdPlugin(version=123)
+
+
+@pytest.mark.asyncio
+async def test_force_new_uuid(
+    plugin: plugins.CorrelationIdPlugin,
+    mocked_request: Request,
+    mocked_response: Response,
+):
+    plugin.force_new_uuid = True
+    await plugin.process_request(mocked_request)
+    await plugin.enrich_response(mocked_response)
+
+    assert (
+        dummy_correlation_id
+        != mocked_response.headers[HeaderKeys.correlation_id]
+    )
+
+
+@pytest.mark.asyncio
+async def test_uuid_validation(
+    plugin: plugins.CorrelationIdPlugin,
+    mocked_request: Request,
+    mocked_response: Response,
+):
+    mocked_request.headers[HeaderKeys.correlation_id] = "invalid_uuid"
+    with pytest.raises(ValueError):
+        await plugin.process_request(mocked_request)

tests/test_unit/test_plugins/test_request_id.py

@@ -49,3 +49,12 @@ async def test_enrich_response_str(
     await plugin.enrich_response(mocked_response)
 
     assert dummy_request_id == mocked_response.headers[HeaderKeys.request_id]
+
+
+@pytest.mark.asyncio
+async def test_invalid_request_id_uuid(
+    plugin: plugins.RequestIdPlugin, mocked_request: Request
+):
+    mocked_request.headers[HeaderKeys.request_id] = "invalid_uuid"
+    with pytest.raises(ValueError):
+        await plugin.process_request(mocked_request)