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 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}"
    )