cleveragents-core/features/steps/a2a_sse_streaming_steps.py

"""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 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: Any, event_type: str, plan_id: str) -> None:
    context.event = A2aEvent(
        event_type=event_type,
        plan_id=plan_id,
        data={"status": "working"},
    )


@when("I format the event as SSE")
def step_format_sse(context: Any) -> None:
    context.sse_output = SseEventFormatter.format(context.event)


@then('the SSE output should contain "{text}"')
def step_sse_contains(context: Any, 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: Any) -> 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: Any) -> None:
    context.keepalive = SseEventFormatter.format_keepalive()


@then('the keepalive should start with ":"')
def step_keepalive_starts_colon(context: Any) -> 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: Any) -> None:
    context.queue = A2aEventQueue()


@given("a mock EventBus")
def step_create_mock_bus(context: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> 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: Any) -> None:
    context.queue.close()


@then("the queue should contain a TaskStatusUpdateEvent")
def step_queue_has_status_event(context: Any) -> 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: Any) -> 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: Any) -> 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: Any, value: str) -> None:
    assert value == TASK_STATUS_UPDATE


@then('the TASK_ARTIFACT_UPDATE constant should equal "{value}"')
def step_check_artifact_constant(context: Any, value: str) -> None:
    assert value == TASK_ARTIFACT_UPDATE


@then("the SSE data line should contain valid JSON")
def step_sse_data_valid_json(context: Any) -> 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: Any, key: str, value: str) -> None:
    assert context.sse_json.get(key) == value, (
        f"Expected JSON['{key}'] == '{value}', got: {context.sse_json.get(key)!r}"
    )