import pytest

from gaia_agent.observability import traced_step


class FakeSpan:
    def __init__(self):
        self.updates = []
        self.ended = False

    def update(self, **kwargs):
        self.updates.append(kwargs)

    def end(self):
        self.ended = True


class FakeTrace:
    def __init__(self):
        self.span = FakeSpan()

    def start_observation(self, name, as_type):
        self.name = name
        self.as_type = as_type
        return self.span


def test_traced_step_records_output_then_ends_span():
    trace = FakeTrace()

    output = traced_step(trace, "step", lambda: {"answer": "Paris"})

    assert output == {"answer": "Paris"}
    assert trace.name == "step"
    assert trace.as_type == "span"
    assert trace.span.updates == [{"output": {"answer": "Paris"}}]
    assert trace.span.ended is True


def test_traced_step_records_error_then_reraises():
    trace = FakeTrace()

    with pytest.raises(ValueError):
        traced_step(trace, "step", lambda: (_ for _ in ()).throw(ValueError("bad")))

    assert trace.span.updates == [{"level": "ERROR", "status_message": "bad"}]
    assert trace.span.ended is True