diff --git a/ocpp/charge_point.py b/ocpp/charge_point.py index 2b9da5f24..57caefea6 100644 --- a/ocpp/charge_point.py +++ b/ocpp/charge_point.py @@ -435,9 +435,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): diff --git a/ocpp/routing.py b/ocpp/routing.py index bf4b12f2b..d9b4c311f 100644 --- a/ocpp/routing.py +++ b/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 diff --git a/tests/test_charge_point.py b/tests/test_charge_point.py index 25a7c8e0d..eee073e1a 100644 --- a/tests/test_charge_point.py +++ b/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