Feat/allow passing answer from @on to @after hook

ocpp/charge_point.py

@@ -402,9 +402,12 @@ async def _handle_call(self, msg):
             # call_unique_id should be passed as kwarg only if is defined explicitly
             # in the handler signature
             if call_unique_id_required:
-                response = handler(**snake_case_payload, call_unique_id=msg.unique_id)
-            else:
-                response = handler(**snake_case_payload)
+                snake_case_payload["call_unique_id"] = msg.unique_id
+            # call_response should be passed as kwarg only if the after handler
+            # was decorated with inject_response=True
+            if getattr(handler, "_inject_response", False):
+                snake_case_payload["call_response"] = response_payload
+            response = handler(**snake_case_payload)
             # Create task to avoid blocking when making a call inside the
             # after handler
             if inspect.isawaitable(response):

ocpp/routing.py

@@ -56,7 +56,7 @@ def inner(*args, **kwargs):
     return decorator
 
 
-def after(action):
+def after(action, inject_response=False):
     """Function decorator to mark function as hook to post-request hook.
 
     This hook's arguments are the data that is in the payload for the specific
@@ -68,6 +68,17 @@ def after(action):
         def after_boot_notification():
             pass
 
+    When ``inject_response`` is set to ``True``, the response that was returned
+    by the matching ``@on`` handler (and sent back to the counterparty) is
+    passed to the hook as the ``call_response`` keyword argument. This avoids
+    having to store the response in a temporary variable to make it available
+    in the ``@after`` hook. It defaults to ``False`` to preserve backwards
+    compatibility:
+
+        @after(Action.boot_notification, inject_response=True):
+        def after_boot_notification(self, call_response, **kwargs):
+            ...
+
     """
 
     def decorator(func):
@@ -76,6 +87,7 @@ def inner(*args, **kwargs):
             return func(*args, **kwargs)
 
         inner._after_action = action
+        inner._inject_response = inject_response
         if func.__name__ not in routables:
             routables.append(func.__name__)
         return inner

tests/test_charge_point.py

@@ -476,3 +476,38 @@ def after_boot_notification(self, *args, **kwargs):
     assert ChargerA.after_boot_notification_call_count == 1
     assert ChargerB.on_boot_notification_call_count == 1
     assert ChargerB.after_boot_notification_call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_response_injected_to_after_handler(connection):
+    """
+    This test ensures that the response is injected to the `after` handler
+    when `inject_response` is set to True.
+    """
+
+    class TestChargePoint(cp_16):
+        after_boot_notification_call_count = 0
+
+        @on(Action.boot_notification)
+        def on_boot_notification(self, **kwargs):
+            return BootNotificationResult(
+                current_time="2024-11-01T00:00:00Z",
+                interval=300,
+                status=RegistrationStatus.accepted,
+            )
+
+        @after(Action.boot_notification, inject_response=True)
+        def after_boot_notification(self, call_response, **kwargs):
+            assert call_response["current_time"] == "2024-11-01T00:00:00Z"
+            assert call_response["interval"] == 300
+            assert call_response["status"] == RegistrationStatus.accepted
+            TestChargePoint.after_boot_notification_call_count += 1
+
+    charge_point = TestChargePoint("test_cp", connection)
+    payload = {"chargePointVendor": "vendor", "chargePointModel": "model"}
+    msg = Call(unique_id="1234", action=Action.boot_notification.value, payload=payload)
+    await charge_point._handle_call(msg)
+
+    # Ensure the after handler actually ran, so the assertions above are not
+    # silently skipped.
+    assert TestChargePoint.after_boot_notification_call_count == 1