"""Step definitions for a2a_sse_streaming.feature.""" from __future__ import annotations import json from typing import Any from unittest.mock import MagicMock from behave import given, then, when from behave.runner import Context from cleveragents.a2a.events import ( TASK_ARTIFACT_UPDATE, TASK_STATUS_UPDATE, A2aEventQueue, EventBusBridge, SseEventFormatter, ) from cleveragents.a2a.models import A2aEvent @given('an A2aEvent with type "{event_type}" and plan_id "{plan_id}"') def step_create_event(context: Context, event_type: str, plan_id: str) -> None: context.event = A2aEvent( event_type=event_type, plan_id=plan_id, data={"status": "working"}, ) @given( 'an A2aEvent with type "{event_type}" and plan_id "{plan_id}" and data {data_json}' ) def step_create_event_with_data( context: Context, event_type: str, plan_id: str, data_json: str ) -> None: data: dict[str, Any] = json.loads(data_json) context.event = A2aEvent( event_type=event_type, plan_id=plan_id, data=data, ) @given('an A2aEvent with type "{event_type}" and no plan_id') def step_create_event_no_plan_id(context: Context, event_type: str) -> None: context.event = A2aEvent( event_type=event_type, plan_id=None, data={"status": "working"}, ) @when("I format the event as SSE") def step_format_sse(context: Context) -> None: context.sse_output = SseEventFormatter.format(context.event) @then('the SSE output should contain "{text}"') def step_sse_contains(context: Context, text: str) -> None: assert text in context.sse_output, ( f"Expected '{text}' in SSE output: {context.sse_output!r}" ) @then("the SSE output should end with two newlines") def step_sse_ends_newlines(context: Context) -> None: assert context.sse_output.endswith("\n\n"), ( f"SSE output should end with two newlines: {context.sse_output!r}" ) @when("I format a keepalive SSE message") def step_format_keepalive(context: Context) -> None: context.keepalive = SseEventFormatter.format_keepalive() @then('the keepalive should start with ":"') def step_keepalive_starts_colon(context: Context) -> None: assert context.keepalive.startswith(":"), ( f"Keepalive should start with ':', got: {context.keepalive!r}" ) @given("an A2aEventQueue for SSE testing") def step_create_queue(context: Context) -> None: context.queue = A2aEventQueue() @given("a mock EventBus") def step_create_mock_bus(context: Context) -> None: context.bus = MagicMock() context.bus_callback = None def fake_subscribe(callback: Any) -> MagicMock: context.bus_callback = callback return MagicMock() context.bus.subscribe.side_effect = fake_subscribe @given("an EventBusBridge connecting bus to queue") def step_create_bridge(context: Context) -> None: context.bridge = EventBusBridge(context.bus, context.queue) context.bridge.start() @when("the bridge receives a PLAN_CREATED domain event") def step_bridge_plan_created(context: Context) -> None: domain_event = MagicMock() domain_event.event_type = MagicMock(value="PLAN_CREATED") domain_event.plan_id = "plan-test-001" domain_event.details = {"action": "test"} context.bridge_error = None try: context.bus_callback(domain_event) except Exception as e: context.bridge_error = e @when("the bridge receives a CHECKPOINT_RESTORED domain event") def step_bridge_checkpoint_restored(context: Context) -> None: domain_event = MagicMock() domain_event.event_type = MagicMock(value="CHECKPOINT_RESTORED") domain_event.plan_id = "plan-test-002" domain_event.details = {"checkpoint_id": "cp-001"} context.bus_callback(domain_event) @when("the queue is closed") def step_close_queue(context: Context) -> None: context.queue.close() @then("the queue should contain a TaskStatusUpdateEvent") def step_queue_has_status_event(context: Context) -> None: events = context.queue.get_events() types = [e.event_type for e in events] assert TASK_STATUS_UPDATE in types, ( f"Expected {TASK_STATUS_UPDATE} in queue, got: {types}" ) @then("the queue should contain a TaskArtifactUpdateEvent") def step_queue_has_artifact_event(context: Context) -> None: events = context.queue.get_events() types = [e.event_type for e in events] assert TASK_ARTIFACT_UPDATE in types, ( f"Expected {TASK_ARTIFACT_UPDATE} in queue, got: {types}" ) @then("no error should be raised from the bridge") def step_no_bridge_error(context: Context) -> None: assert context.bridge_error is None, ( f"Expected no error, got: {context.bridge_error}" ) @then('the TASK_STATUS_UPDATE constant should equal "{value}"') def step_check_status_constant(context: Context, value: str) -> None: assert value == TASK_STATUS_UPDATE @then('the TASK_ARTIFACT_UPDATE constant should equal "{value}"') def step_check_artifact_constant(context: Context, value: str) -> None: assert value == TASK_ARTIFACT_UPDATE @then("the SSE data line should contain valid JSON") def step_sse_data_valid_json(context: Context) -> None: for line in context.sse_output.splitlines(): if line.startswith("data: "): payload = line[6:] context.sse_json = json.loads(payload) return raise AssertionError("No data: line found in SSE output") @then('the JSON should have key "{key}" with value "{value}"') def step_json_has_key_value(context: Context, key: str, value: str) -> None: assert context.sse_json.get(key) == value, ( f"Expected JSON['{key}'] == '{value}', got: {context.sse_json.get(key)!r}" ) @then('the JSON should have a "{key}" object') def step_json_has_object(context: Context, key: str) -> None: assert key in context.sse_json, ( f"Expected JSON to have key '{key}', got keys: {list(context.sse_json.keys())}" ) assert isinstance(context.sse_json[key], dict), ( f"Expected JSON['{key}'] to be a dict, got: {type(context.sse_json[key])!r}" ) context.sse_params = context.sse_json[key] @then('the JSON params should have key "{key}" with value "{value}"') def step_json_params_has_key_value(context: Context, key: str, value: str) -> None: params: dict[str, Any] = context.sse_json.get("params", {}) assert params.get(key) == value, ( f"Expected params['{key}'] == '{value}', got: {params.get(key)!r}" ) @then('the JSON params should have key "{key}" with numeric value {value:g}') def step_json_params_has_numeric_value( context: Context, key: str, value: float ) -> None: params: dict[str, Any] = context.sse_json.get("params", {}) actual = params.get(key) assert actual == value, f"Expected params['{key}'] == {value}, got: {actual!r}" @then('the JSON params should not have key "{key}"') def step_json_params_not_have_key(context: Context, key: str) -> None: params: dict[str, Any] = context.sse_json.get("params", {}) assert key not in params, ( f"Expected params to NOT have key '{key}', but found: {params.get(key)!r}" ) @then('the JSON should not have key "{key}"') def step_json_not_have_key(context: Context, key: str) -> None: assert key not in context.sse_json, ( f"Expected JSON to NOT have key '{key}', but found: {context.sse_json.get(key)!r}" )