microsoft--agent-lightning/tests/tracer/test_weave.py

# Copyright (c) Microsoft. All rights reserved.

from __future__ import annotations

import datetime
from types import SimpleNamespace

import pytest

from agentlightning.store.base import LightningStore
from agentlightning.tracer.weave import WeaveTracer
from agentlightning.types import Span


class MockLightningStore(LightningStore):
    """A minimal stub-only LightningStore, only implements methods likely used in tests."""

    def __init__(self) -> None:
        super().__init__()
        self.spans: list[Span] = []

    async def get_next_span_sequence_id(self, rollout_id: str, attempt_id: str) -> int:
        return len(self.spans)

    async def add_span(self, span: Span) -> Span:
        self.spans.append(span)
        return span

    def clear_spans(self) -> None:
        self.spans = []

    def get_traces(self) -> list[Span]:
        return self.spans


class FakeWeaveCall:
    def __init__(
        self,
        op: str,
        attributes: dict[str, object] | None = None,
        inputs: dict[str, object] | None = None,
    ):
        base_time = datetime.datetime.fromtimestamp(96400, tz=datetime.timezone.utc)
        self.op_name = op
        self.attributes = attributes or {}
        self.inputs = inputs or {}
        self.output = None
        self.summary: dict[str, object] | None = {}
        self.started_at = base_time
        self.ended_at = base_time
        self.exception: str | None = None


class FakeWeaveClient:
    def __init__(self, server: object):
        self.server = server
        self.project = "fake-project"
        self.created_calls: list[FakeWeaveCall] = []
        self.finished_calls: list[FakeWeaveCall] = []

    def create_call(
        self,
        *,
        op: str,
        attributes: dict[str, object] | None = None,
        inputs: dict[str, object] | None = None,
    ) -> FakeWeaveCall:
        call = FakeWeaveCall(op=op, attributes=attributes or {}, inputs=inputs or {})
        self.created_calls.append(call)
        return call

    def finish_call(self, call: FakeWeaveCall, exception: Exception | None = None) -> None:
        if exception is not None:
            call.exception = str(exception)
        call.ended_at = call.started_at + datetime.timedelta(seconds=1)
        self.finished_calls.append(call)

    def flush(self) -> None:
        pass


def _func_without_exception():
    """Function that always executed successfully to test success tracing."""
    pass


def _func_with_exception():
    """Function that always executed successfully to test success tracing."""
    raise ValueError("This is a test exception")


@pytest.mark.weave
def test_weave_tracer_create_span(monkeypatch: pytest.MonkeyPatch) -> None:
    import weave  # type: ignore

    tracer = WeaveTracer(instrument_managed=False)
    fake_client = FakeWeaveClient(server=tracer._server)  # pyright: ignore[reportPrivateUsage]
    monkeypatch.setattr(weave, "get_client", lambda: fake_client)

    span = tracer.create_span("weave-span", attributes={"foo": "bar"})

    assert span.name == "weave-span"
    assert span.attributes == {"foo": "bar"}
    assert len(fake_client.created_calls) == 1
    assert len(fake_client.finished_calls) == 1


@pytest.mark.weave
def test_weave_tracer_operation_context_records_exception(monkeypatch: pytest.MonkeyPatch) -> None:
    import weave  # type: ignore

    tracer = WeaveTracer(instrument_managed=False)
    fake_client = FakeWeaveClient(server=tracer._server)  # pyright: ignore[reportPrivateUsage]
    monkeypatch.setattr(weave, "get_client", lambda: fake_client)

    with pytest.raises(RuntimeError):
        with tracer.operation_context("weave-op", attributes={"foo": "bar"}) as ctx:
            raise RuntimeError("failed")

    recorded_span = ctx.get_recorded_span()  # type: ignore
    assert recorded_span.name == "weave-op"
    assert recorded_span.status.status_code == "ERROR"
    assert recorded_span.status.description == "failed"
    assert len(fake_client.finished_calls) == 1


@pytest.mark.weave
@pytest.mark.asyncio
async def test_weave_trace_without_exception():
    tracer = WeaveTracer()
    store = MockLightningStore()
    tracer.init()
    tracer.init_worker(0, store=store)

    try:
        # Case where store, rollout_id, and attempt_id are all non-none.
        async with tracer.trace_context(name="weave_test", rollout_id="test_rollout_id", attempt_id="test_attempt_id"):
            _func_without_exception()
        print(tracer.get_last_trace())
        spans = store.get_traces()
        assert len(spans) > 0

        has_error = False
        for span in spans:
            has_error = getattr(span, "status", None) and span.status.status_code == "ERROR"

        assert not has_error
    finally:
        tracer.teardown_worker(0)
        tracer.teardown()


@pytest.mark.weave
@pytest.mark.asyncio
async def test_weave_trace_with_exception():
    tracer = WeaveTracer()
    store = MockLightningStore()
    tracer.init()
    tracer.init_worker(0, store=store)

    try:
        with pytest.raises(ValueError):
            # Case where store, rollout_id, and attempt_id are all non-none.
            async with tracer.trace_context(
                name="weave_test", rollout_id="test_rollout_id", attempt_id="test_attempt_id"
            ):
                _func_with_exception()
        spans = store.get_traces()
        assert len(spans) > 0

        has_error = False
        for span in spans:
            has_error = getattr(span, "status", None) and span.status.status_code == "ERROR"

        assert has_error
    finally:
        tracer.teardown_worker(0)
        tracer.teardown()


@pytest.mark.weave
@pytest.mark.asyncio
async def test_weave_with_op():
    import weave  # type: ignore

    @weave.op  # type: ignore
    def _func_with_op():
        """Function that always executed successfully to test success tracing."""
        pass

    tracer = WeaveTracer()
    store = MockLightningStore()
    tracer.init()
    tracer.init_worker(0, store=store)

    try:
        # Case where store, rollout_id, and attempt_id are all non-none.
        async with tracer.trace_context(name="weave_test", rollout_id="test_rollout_id", attempt_id="test_attempt_id"):
            _func_with_op()
        spans = store.get_traces()
        len_spans_with_op = len(spans)

        store.clear_spans()
        async with tracer.trace_context(name="weave_test", rollout_id="test_rollout_id", attempt_id="test_attempt_id"):
            _func_without_exception()
        spans = store.get_traces()
        len_spans_without_op = len(spans)

        assert len_spans_with_op > 0
        assert len_spans_without_op > 0
        assert len_spans_with_op > len_spans_without_op

    finally:
        tracer.teardown_worker(0)
        tracer.teardown()


@pytest.mark.weave
@pytest.mark.asyncio
async def test_weave_trace_call_to_span():
    child = SimpleNamespace(
        op_name="child_func",
        inputs={"child_input": "x"},
        output={"child_output": 42},
        attributes={"child_attribute": "z"},
        summary={"status_counts": {"success": 1, "error": 0}},
        _children=[],
        started_at=datetime.datetime(2025, 12, 1, 0, 0, 1, tzinfo=datetime.timezone.utc),
        ended_at=datetime.datetime(2025, 12, 1, 0, 0, 2, tzinfo=datetime.timezone.utc),
        trace_id="trace-1",
        id="span-2",
        parent_id="span-1",
        func_name="child_func",
        exception=None,
    )

    parent = SimpleNamespace(
        op_name="parent_func",
        inputs={"parent_input": "y"},
        output={"parent_output": 99},
        attributes={"parent_attribute": "y"},
        summary={"status_counts": {"success": 1, "error": 0}},
        _children=[child],
        started_at=datetime.datetime(2025, 12, 1, 0, 0, 0, tzinfo=datetime.timezone.utc),
        ended_at=datetime.datetime(2025, 12, 1, 0, 0, 1, tzinfo=datetime.timezone.utc),
        trace_id="trace-1",
        id="span-1",
        parent_id=None,
        func_name="parent_func",
        exception=None,
    )

    tracer = WeaveTracer()
    parent_span = await tracer.convert_call_to_span(parent)  # type: ignore
    assert parent_span.attributes["agentlightning.operation.input.parent_input"] == "y"
    assert parent_span.attributes["agentlightning.operation.output.parent_output"] == 99

    child_span = await tracer.convert_call_to_span(child)  # type: ignore
    assert child_span.attributes["agentlightning.operation.input.child_input"] == "x"
    assert child_span.attributes["agentlightning.operation.output.child_output"] == 42
    assert child_span.parent_id == "span-1"