Add default custom header support to client

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "toadr3"
-version = "0.25.0"
+version = "0.26.0"
 description = "Tiny OpenADR 3 compatible client Python Library"
 authors = ["Jean-Paul Balabanian <jean-paul.balabanian@eviny.no>"]
 license = "Apache-2.0"

tests/_reports_response.py

@@ -57,15 +57,20 @@ async def reports_get_response(request: web.Request) -> web.Response:
 async def reports_post_response(request: web.Request) -> web.Response:
     auth = request.headers.get("Authorization", None)
 
-    if auth is None:
-        data = {
-            "status": 403,
-            "title": "Forbidden",
-        }
-        return web.json_response(data=data, status=403)
+    credential_response = check_credentials(auth)
+    if credential_response is not None:
+        return credential_response
 
-    report_data = await request.json()
+    # If custom header is set but not set to "CustomValue" return 400
     custom_header = request.headers.get("X-Custom-Header", None)
+    extra_header_response = check_custom_header(custom_header)
+    if extra_header_response is not None:
+        return extra_header_response
+
+    report_data = await request.json()
+
+    if custom_header == "CustomValue":
+        report_data["clientName"] = "CustomClientName"
 
     if report_data["clientName"] == "Unauthorized":
         data = {
@@ -83,21 +88,6 @@ async def reports_post_response(request: web.Request) -> web.Response:
         }
         return web.json_response(data=data, status=409)
 
-    # If custom header is not set to "CustomValue" return 400 or update clientName
-    if custom_header is not None:
-        match custom_header:
-            case "CustomValue":
-                report_data["clientName"] = "CustomClientName"
-            case _:
-                return web.json_response(
-                    data={
-                        "status": 400,
-                        "title": "Bad Request",
-                        "detail": f"Invalid value for X-Custom-Header: {custom_header}",
-                    },
-                    status=400,
-                )
-
     # Return the report data with some additional fields
     report_data["id"] = "1234"
     report_data["createdDateTime"] = "2024-09-30T12:12:34Z"

tests/_subscriptions_response.py

@@ -76,9 +76,17 @@ async def subscriptions_post_response(request: web.Request) -> web.Response:
         }
         return web.json_response(data=data, status=403)
 
-    subscription_data = await request.json()
     custom_header = request.headers.get("X-Custom-Header", None)
 
+    extra_header_response = check_custom_header(custom_header)
+    if extra_header_response is not None:
+        return extra_header_response
+
+    subscription_data = await request.json()
+
+    if custom_header == "CustomValue":
+        subscription_data["clientName"] = "CustomClientName"
+
     if subscription_data["programID"] == "35" or subscription_data["id"] == "123":
         data = {
             "status": 409,
@@ -87,21 +95,6 @@ async def subscriptions_post_response(request: web.Request) -> web.Response:
         }
         return web.json_response(data=data, status=409)
 
-    # If custom header is not set to "CustomValue" return 400
-    if custom_header is not None:
-        match custom_header:
-            case "CustomValue":
-                subscription_data["clientName"] = "CustomClientName"
-            case _:
-                return web.json_response(
-                    data={
-                        "status": 400,
-                        "title": "Bad Request",
-                        "detail": f"Invalid value for X-Custom-Header: {custom_header}",
-                    },
-                    status=400,
-                )
-
     # Return the report data with some additional fields
     subscription_data["id"] = "123"
     subscription_data["createdDateTime"] = "2024-09-30T12:12:34Z"

tests/test_client.py

@@ -1,4 +1,6 @@
 import pytest
+from pydantic import BaseModel
+from testdata import default_report_model, default_subscription_model
 
 import toadr3
 
@@ -14,3 +16,34 @@ async def test_client_context_manager(client: toadr3.ToadrClient) -> None:
     with pytest.raises(RuntimeError, match="Client is closed"):
         async with client:
             pass
+
+
+@pytest.mark.parametrize(
+    ("method_name", "method_args"),
+    [
+        ("get_programs", ()),
+        ("get_events", ()),
+        ("get_reports", ()),
+        ("post_report", (default_report_model(),)),
+        ("get_subscriptions", ()),
+        ("post_subscription", (default_subscription_model(),)),
+    ],
+)
+async def test_client_default_custom_headers(
+    client: toadr3.ToadrClient, method_name: str, method_args: tuple[BaseModel]
+) -> None:
+    # Test that default_custom_headers are set and forwarded correctly
+    assert client.default_custom_headers == {}
+
+    custom_headers = {"X-Custom-Header": "InvalidValue"}
+    client_with_headers = toadr3.ToadrClient(
+        vtn_url="vtn_url",
+        oauth_config=client._oauth_config,  # noqa: SLF001
+        session=client.client_session,
+        default_custom_headers=custom_headers,
+    )
+    assert client_with_headers.default_custom_headers == custom_headers
+
+    msg = "Bad Request 400 - Invalid value for X-Custom-Header: InvalidValue"
+    with pytest.raises(toadr3.ToadrError, match=msg):
+        _ = await getattr(client_with_headers, method_name)(*method_args)

tests/test_json_encoder.py

@@ -3,7 +3,14 @@
 
 import pytest
 from pydantic import BaseModel
-from testdata import create_event, create_report
+from testdata import (
+    create_event,
+    create_report,
+    default_event,
+    default_program,
+    default_report,
+    default_subscription,
+)
 
 from toadr3.models import Event, Program, Report, Subscription
 
@@ -43,57 +50,6 @@ def test_json_encoder_event() -> None:
     assert result == data
 
 
-def default_event() -> dict[str, Any]:
-    return {
-        "programID": "11",
-        "intervals": [
-            {
-                "id": 1,
-                "payloads": [{"type": "type", "values": ["value1"]}],
-            }
-        ],
-    }
-
-
-def default_report() -> dict[str, Any]:
-    return {
-        "programID": "11",
-        "eventID": "33",
-        "clientName": "YAC",
-        "resources": [
-            {
-                "resourceName": "resource1",
-                "intervals": [
-                    {
-                        "id": 1,
-                        "payloads": [{"type": "type", "values": ["value1"]}],
-                    }
-                ],
-            }
-        ],
-    }
-
-
-def default_program() -> dict[str, Any]:
-    return {
-        "programName": "pname",
-    }
-
-
-def default_subscription() -> dict[str, Any]:
-    return {
-        "clientName": "YAC",
-        "programID": "11",
-        "objectOperations": [
-            {
-                "objects": ["EVENT"],
-                "operations": ["POST"],
-                "callbackUrl": "https://example.com/callback",
-            }
-        ],
-    }
-
-
 @pytest.mark.parametrize(
     ("model_class", "model_data"),
     [

tests/testdata.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from toadr3.models import ObjectType, OperationType
+from toadr3.models import Event, ObjectType, OperationType, Program, Report, Subscription
 
 
 def create_event(**kwargs: str | int) -> dict[str, Any]:
@@ -171,3 +171,71 @@ def create_subscriptions() -> list[dict[str, Any]]:
         create_subscription("4", "5"),
         exception,
     ]
+
+
+# The following are default data dictionaries for creating minimal models
+def default_event() -> dict[str, Any]:
+    return {
+        "programID": "11",
+        "intervals": [
+            {
+                "id": 1,
+                "payloads": [{"type": "type", "values": ["value1"]}],
+            }
+        ],
+    }
+
+
+def default_event_model() -> Event:
+    return Event.model_validate(default_event())
+
+
+def default_report() -> dict[str, Any]:
+    return {
+        "programID": "11",
+        "eventID": "33",
+        "clientName": "YAC",
+        "resources": [
+            {
+                "resourceName": "resource1",
+                "intervals": [
+                    {
+                        "id": 1,
+                        "payloads": [{"type": "type", "values": ["value1"]}],
+                    }
+                ],
+            }
+        ],
+    }
+
+
+def default_report_model() -> Report:
+    return Report.model_validate(default_report())
+
+
+def default_program() -> dict[str, Any]:
+    return {
+        "programName": "pname",
+    }
+
+
+def default_program_model() -> Program:
+    return Program.model_validate(default_program())
+
+
+def default_subscription() -> dict[str, Any]:
+    return {
+        "clientName": "YAC",
+        "programID": "11",
+        "objectOperations": [
+            {
+                "objects": ["EVENT"],
+                "operations": ["POST"],
+                "callbackUrl": "https://example.com/callback",
+            }
+        ],
+    }
+
+
+def default_subscription_model() -> Subscription:
+    return Subscription.model_validate(default_subscription())

toadr3/client.py

@@ -20,6 +20,7 @@ def __init__(
         vtn_url: str,
         oauth_config: toadr3.OAuthConfig | None,
         session: ClientSession | None = None,
+        default_custom_headers: dict[str, str] | None = None,
     ) -> None:
         """Initialize the client.
 
@@ -31,7 +32,10 @@ def __init__(
             The OAuth configuration to use or None if credentials are not required.
         session : ClientSession | None
             The session to use or None if the client should make its own session.
+        default_custom_headers : dict[str, str] | None
+            Default custom headers to include in every request.
         """
+        self._default_custom_headers = default_custom_headers or {}
         self._vtn_url = vtn_url.rstrip("/")
         self._oauth_config = oauth_config
         self._session = session if session is not None else ClientSession()
@@ -54,6 +58,11 @@ def closed(self) -> bool:
         """Whether or not the client is closed."""
         return self._closed
 
+    @property
+    def default_custom_headers(self) -> dict[str, str]:
+        """Default custom headers to include in every request."""
+        return self._default_custom_headers
+
     async def _fetch_token(self) -> toadr3.AccessToken | None:
         """Fetch an access token from the OAuth2 server."""
         if self._oauth_config is None:
@@ -69,6 +78,11 @@ async def token(self) -> toadr3.AccessToken | None:
 
             return self._token
 
+    def _prepare_headers(self, custom_headers: dict[str, str] | None) -> dict[str, str]:
+        """Prepare headers for the request."""
+        custom_headers = custom_headers or {}
+        return self._default_custom_headers | custom_headers
+
     async def get_events(
         self,
         program_id: str | None = None,
@@ -125,7 +139,7 @@ async def get_events(
             skip=skip,
             limit=limit,
             extra_params=extra_params,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def get_programs(
@@ -180,7 +194,7 @@ async def get_programs(
             skip=skip,
             limit=limit,
             extra_params=extra_params,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def get_subscriptions(
@@ -248,7 +262,7 @@ async def get_subscriptions(
             skip=skip,
             limit=limit,
             extra_params=extra_params,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def post_subscription(
@@ -285,7 +299,7 @@ async def post_subscription(
             vtn_url=self._vtn_url,
             access_token=await self.token,
             subscription=subscription,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def get_reports(
@@ -341,7 +355,7 @@ async def get_reports(
             skip=skip,
             limit=limit,
             extra_params=extra_params,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def post_report(
@@ -377,7 +391,7 @@ async def post_report(
             vtn_url=self._vtn_url,
             access_token=await self.token,
             report=report,
-            custom_headers=custom_headers,
+            custom_headers=self._prepare_headers(custom_headers),
         )
 
     async def close(self) -> None: