spanforge 2.0.0__py3-none-any.whl

spanforge/testing.py

@@ -0,0 +1,376 @@
+"""spanforge.testing — Test utilities for SpanForge SDK consumers.
+
+Provides helpers for writing unit and integration tests that involve
+SpanForge events, exporters, and the trace store.  Designed to be imported
+only in test code (not in production).
+
+Usage::
+
+    from spanforge.testing import capture_events, MockExporter, assert_event_schema_valid
+
+    # Capture all events emitted during a block
+    with capture_events() as captured:
+        # code that emits events
+        ...
+
+    assert len(captured) == 1
+    assert captured[0].event_type == "llm.trace.span.completed"
+    assert_event_schema_valid(captured[0])
+
+    # Inject a mock exporter
+    mock = MockExporter()
+    with mock.installed():
+        # code that emits events
+        ...
+    assert len(mock.events) == 1
+
+    # Isolated TraceStore for one test
+    from spanforge.testing import trace_store
+    with trace_store() as store:
+        configure(enable_trace_store=True)
+        # ... emit events ...
+        events = store.get_trace(trace_id)
+"""
+
+from __future__ import annotations
+
+import contextlib
+import threading
+from typing import TYPE_CHECKING, Any, Generator
+
+if TYPE_CHECKING:
+    from spanforge.event import Event
+    from spanforge._store import TraceStore
+    from spanforge._span import Span
+
+__all__ = [
+    "MockExporter",
+    "assert_event_schema_valid",
+    "assert_span_emitted",
+    "capture_events",
+    "captured_spans",
+    "trace_store",
+]
+
+
+# ---------------------------------------------------------------------------
+# MockExporter
+# ---------------------------------------------------------------------------
+
+
+class MockExporter:
+    """A synchronous in-memory exporter for testing.
+
+    Records every event passed to :meth:`export` into :attr:`events`.
+    Supports optional filtering, ordered access, and a context manager
+    that temporarily replaces the global exporter.
+
+    Args:
+        raise_on_export: When set to an :class:`Exception` subclass or
+                         instance, :meth:`export` raises it to simulate
+                         export failures.
+
+    Attributes:
+        events: List of all :class:`~spanforge.event.Event` objects exported
+                in chronological order.
+
+    Example::
+
+        mock = MockExporter()
+        with mock.installed():
+            tracer.span("test").__enter__().__exit__(None, None, None)
+
+        assert mock.events[0].event_type == "llm.trace.span.completed"
+    """
+
+    def __init__(
+        self,
+        raise_on_export: type[Exception] | Exception | None = None,
+    ) -> None:
+        self.events: list[Event] = []
+        self._lock = threading.Lock()
+        self._raise_on_export = raise_on_export
+
+    def export(self, event: Event) -> None:
+        """Record *event*.  Raises configured exception if one is set.
+
+        Args:
+            event: The event to record.
+
+        Raises:
+            Exception: The configured ``raise_on_export`` exception, if set.
+        """
+        if self._raise_on_export is not None:
+            if isinstance(self._raise_on_export, type):
+                raise self._raise_on_export("MockExporter.raise_on_export triggered")
+            raise self._raise_on_export
+        with self._lock:
+            self.events.append(event)
+
+    async def export_batch(self, events: Any) -> None:  # NOSONAR
+        """Async batch export — records all events in *events*.
+
+        Args:
+            events: Iterable of :class:`~spanforge.event.Event` objects.
+        """
+        for event in events:
+            self.export(event)
+
+    def clear(self) -> None:
+        """Remove all recorded events."""
+        with self._lock:
+            self.events.clear()
+
+    def filter_by_type(self, event_type: str) -> list[Event]:
+        """Return all recorded events matching *event_type*.
+
+        Args:
+            event_type: Dotted event type string, e.g.
+                        ``"llm.trace.span.completed"``.
+
+        Returns:
+            Filtered, ordered list of matching events.
+        """
+        et = str(event_type)
+        with self._lock:
+            return [
+                e for e in self.events
+                if (e.event_type.value if hasattr(e.event_type, "value") else str(e.event_type)) == et
+            ]
+
+    @contextlib.contextmanager
+    def installed(self) -> Generator[MockExporter, None, None]:
+        """Context manager that installs this exporter as the global exporter.
+
+        Replaces the SDK's active exporter for the duration of the block,
+        then restores the original state::
+
+            mock = MockExporter()
+            with mock.installed():
+                ...  # all events go to mock.events
+
+        Yields:
+            This :class:`MockExporter` instance.
+        """
+        from spanforge._stream import _exporter_lock  # noqa: PLC0415
+        import spanforge._stream as _stream  # noqa: PLC0415
+
+        # Save state
+        with _exporter_lock:
+            original = _stream._cached_exporter
+            _stream._cached_exporter = self
+        try:
+            yield self
+        finally:
+            with _exporter_lock:
+                _stream._cached_exporter = original
+
+    def __repr__(self) -> str:
+        return f"MockExporter(events={len(self.events)})"
+
+
+# ---------------------------------------------------------------------------
+# capture_events()
+# ---------------------------------------------------------------------------
+
+
+@contextlib.contextmanager
+def capture_events() -> Generator[list[Event], None, None]:
+    """Context manager that captures all events emitted during the block.
+
+    Events are collected into a list that is yielded and grows in real-time
+    as events are emitted.  The original exporter is restored on exit.
+
+    Example::
+
+        with capture_events() as events:
+            with tracer.span("test"):
+                pass
+
+        assert events[0].payload["span_name"] == "test"
+
+    Yields:
+        A live ``list[Event]`` that is populated as events are emitted.
+    """
+    mock = MockExporter()
+    with mock.installed():
+        yield mock.events
+
+
+# ---------------------------------------------------------------------------
+# assert_event_schema_valid()
+# ---------------------------------------------------------------------------
+
+
+def assert_event_schema_valid(event: Event) -> None:
+    """Assert that *event* passes SDK schema validation.
+
+    Calls :func:`~spanforge.validate.validate_event` and re-raises any
+    :class:`~spanforge.exceptions.SchemaValidationError` as an
+    :class:`AssertionError` with the original message — so failures
+    surface cleanly in :func:`pytest.raises` and ``assert`` blocks.
+
+    Args:
+        event: The event to validate.
+
+    Raises:
+        AssertionError: If *event* fails schema validation.
+
+    Example::
+
+        from spanforge.testing import assert_event_schema_valid
+        assert_event_schema_valid(my_event)
+    """
+    from spanforge.exceptions import SchemaValidationError  # noqa: PLC0415
+    from spanforge.validate import validate_event  # noqa: PLC0415
+
+    try:
+        validate_event(event)
+    except SchemaValidationError as exc:
+        raise AssertionError(f"Event failed schema validation: {exc}") from exc
+
+
+# ---------------------------------------------------------------------------
+# trace_store() context manager
+# ---------------------------------------------------------------------------
+
+
+@contextlib.contextmanager
+def trace_store(max_traces: int = 100) -> Generator[TraceStore, None, None]:
+    """Context manager that provides an isolated :class:`~spanforge._store.TraceStore`.
+
+    Creates a fresh ``TraceStore`` scoped to the block and temporarily
+    installs it as the global store.  The original store is restored on
+    exit, making this safe to use in parallel tests.
+
+    Args:
+        max_traces: Maximum number of traces to retain in the isolated store.
+
+    Yields:
+        A new :class:`~spanforge._store.TraceStore` instance scoped to the
+        ``with`` block.
+
+    Example::
+
+        from spanforge import configure
+        from spanforge.testing import trace_store
+
+        with trace_store() as store:
+            configure(enable_trace_store=True)
+            with tracer.span("test") as s:
+                pass
+            events = store.get_trace(s.trace_id)
+            assert events is not None
+    """
+    from spanforge._store import trace_store as _store_trace_store  # noqa: PLC0415
+
+    with _store_trace_store(max_traces=max_traces) as store:
+        yield store
+
+
+# ---------------------------------------------------------------------------
+# captured_spans() pytest fixture
+# ---------------------------------------------------------------------------
+
+try:
+    import pytest as _pytest
+
+    @_pytest.fixture
+    def captured_spans() -> Generator[list[Span], None, None]:
+        """pytest fixture that captures all :class:`~spanforge._span.Span` objects
+        completed during a test, regardless of operation type.
+
+        Import this fixture in your test module (or ``conftest.py``) to make it
+        available::
+
+            from spanforge.testing import captured_spans  # re-export for pytest
+
+            def test_my_fn(captured_spans):
+                call_my_function()
+                assert any(s.name == "my-step" for s in captured_spans)
+
+        Each test gets an empty list; spans accumulate as the test runs.
+
+        Yields:
+            A live ``list[Span]`` populated as spans are completed.
+        """
+        from spanforge._hooks import hooks  # noqa: PLC0415
+
+        spans: list[Any] = []
+
+        def _cb(span: Any) -> None:
+            spans.append(span)
+
+        hooks.on_span_end(_cb)
+        try:
+            yield spans  # type: ignore[misc]
+        finally:
+            with hooks._lock:
+                try:
+                    hooks._all_end_hooks.remove(_cb)
+                except ValueError:
+                    pass
+
+except ImportError:
+    # pytest not installed — skip fixture definition (production environments).
+    pass
+
+
+# ---------------------------------------------------------------------------
+# assert_span_emitted()
+# ---------------------------------------------------------------------------
+
+
+def assert_span_emitted(
+    spans: list[Any],
+    *,
+    name: str,
+    model: str | None = None,
+    status: str | None = None,
+    operation: str | None = None,
+) -> Any:
+    """Assert that a span matching the given criteria appears in *spans*.
+
+    Typically used with the :func:`captured_spans` fixture.
+
+    Args:
+        spans:     List of :class:`~spanforge._span.Span` objects (from fixture).
+        name:      Required span name to match.
+        model:     When provided, also checks ``span.model == model``.
+        status:    When provided, also checks ``span.status == status``.
+        operation: When provided, also checks ``span.operation == operation``.
+
+    Returns:
+        The first matching :class:`~spanforge._span.Span`.
+
+    Raises:
+        AssertionError: If no span matches all criteria.
+
+    Example::
+
+        def test_llm_call(captured_spans):
+            run_agent()
+            assert_span_emitted(captured_spans, name="llm-call", model="gpt-4o")
+    """
+    for span in spans:
+        if span.name != name:
+            continue
+        if model is not None and span.model != model:
+            continue
+        if status is not None and span.status != status:
+            continue
+        if operation is not None and str(span.operation) != operation:
+            continue
+        return span
+
+    criteria = f"name={name!r}"
+    if model is not None:
+        criteria += f", model={model!r}"
+    if status is not None:
+        criteria += f", status={status!r}"
+    if operation is not None:
+        criteria += f", operation={operation!r}"
+    raise AssertionError(
+        f"No span matching {criteria} found in {len(spans)} captured span(s). "
+        f"Got: {[s.name for s in spans]}"
+    )

spanforge/trace.py

@@ -0,0 +1,199 @@
+"""spanforge.trace — @trace() function decorator and trace engine (Tool 1 / llm-trace).
+
+Single decorator that instruments any Python function as an SpanForge span::
+
+    from spanforge import trace
+
+    @trace(name="search", model="gpt-4o")
+    def call_llm(prompt: str) -> str: ...
+
+    @trace(name="async-step")
+    async def async_step(x: int) -> dict: ...
+
+    @trace(name="web-search", tool=True)
+    def search_web(query: str) -> list[str]: ...
+
+Supports:
+- Sync and async functions/methods
+- Auto-capture of call arguments and return values (opt-in)
+- Parent-child span relationships via contextvars
+- ``tool=True`` to emit spans with ``operation="execute_tool"``
+- Pytest fixture integration via :func:`~spanforge.testing.captured_spans`
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from typing import Any, Callable, TypeVar
+
+from spanforge._span import SpanContextManager
+
+__all__ = ["trace"]
+
+_F = TypeVar("_F", bound=Callable[..., Any])
+
+
+def _safe_repr(value: Any, max_len: int = 200) -> str:
+    """Return a repr of *value* truncated to *max_len* characters."""
+    try:
+        r = repr(value)
+    except Exception:
+        r = "<unrepresentable>"
+    return r[:max_len] + "..." if len(r) > max_len else r
+
+
+class _TraceDecorator:
+    """Wraps a callable so every invocation is recorded as an SpanForge span.
+
+    Created by :func:`trace`; use :func:`trace` rather than instantiating
+    this class directly.
+    """
+
+    def __init__(
+        self,
+        fn: Callable[..., Any],
+        name: str | None,
+        model: str | None,
+        operation: str,
+        tool: bool,
+        capture_args: bool,
+        capture_return: bool,
+        attributes: dict[str, Any] | None,
+    ) -> None:
+        self._fn = fn
+        self._name = name or fn.__qualname__
+        self._model = model
+        self._operation = operation
+        self._tool = tool
+        self._capture_args = capture_args
+        self._capture_return = capture_return
+        self._attributes = dict(attributes or {})
+        # Preserve __name__, __doc__, __module__, etc. on the wrapper.
+        functools.update_wrapper(self, fn)
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _build_attrs(self, args: tuple[Any, ...], kwargs: dict[str, Any]) -> dict[str, Any]:
+        """Build the initial span attributes dict from static attrs + captured args."""
+        attrs: dict[str, Any] = dict(self._attributes)
+        if self._tool:
+            # Mark span so InspectorSession can detect it even without checking operation.
+            attrs["tool"] = True
+        if self._capture_args or self._tool:
+            try:
+                sig = inspect.signature(self._fn)
+                bound = sig.bind(*args, **kwargs)
+                bound.apply_defaults()
+                for k, v in bound.arguments.items():
+                    attrs[f"arg.{k}"] = _safe_repr(v)
+            except (TypeError, ValueError):
+                pass
+        return attrs
+
+    def _record_return(self, span: Any, result: Any) -> None:
+        if self._capture_return or self._tool:
+            span.set_attribute("return_value", _safe_repr(result))
+
+    # ------------------------------------------------------------------
+    # Sync call
+    # ------------------------------------------------------------------
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        attrs = self._build_attrs(args, kwargs)
+        cm = SpanContextManager(
+            name=self._name,
+            model=self._model,
+            operation=self._operation,
+            attributes=attrs,
+        )
+        with cm as span:
+            result = self._fn(*args, **kwargs)
+            self._record_return(span, result)
+        return result
+
+
+def _make_async_wrapper(decorator: _TraceDecorator, fn: Callable[..., Any]) -> Callable[..., Any]:
+    """Return an async wrapper that runs *fn* inside a span context."""
+
+    @functools.wraps(fn)
+    async def _async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        attrs = decorator._build_attrs(args, kwargs)
+        cm = SpanContextManager(
+            name=decorator._name,
+            model=decorator._model,
+            operation=decorator._operation,
+            attributes=attrs,
+        )
+        async with cm as span:
+            result = await fn(*args, **kwargs)
+            decorator._record_return(span, result)
+        return result
+
+    return _async_wrapper
+
+
+def trace(
+    fn: Callable[..., Any] | None = None,
+    *,
+    name: str | None = None,
+    model: str | None = None,
+    operation: str = "chat",
+    tool: bool = False,
+    capture_args: bool = False,
+    capture_return: bool = False,
+    attributes: dict[str, Any] | None = None,
+) -> Any:
+    """Decorator that instruments a function as an SpanForge span.
+
+    Works with or without parentheses::
+
+        @trace
+        def my_fn(): ...
+
+        @trace(name="custom-name", model="gpt-4o")
+        def my_fn(): ...
+
+    Args:
+        fn:             The function to wrap (only when used bare as ``@trace``).
+        name:           Span name.  Defaults to ``fn.__qualname__``.
+        model:          Model identifier string forwarded to ``SpanPayload``.
+        operation:      GenAI operation name (default ``"chat"``).  Any
+                        :class:`~spanforge.namespaces.trace.GenAIOperationName`
+                        value or a custom string.
+        tool:           When ``True``, marks this as a tool call; sets
+                        ``operation="execute_tool"`` regardless of *operation*.
+        capture_args:   When ``True``, records call arguments as ``arg.<name>``
+                        span attributes (values truncated to 200 chars).
+        capture_return: When ``True``, records the return value as a
+                        ``return_value`` span attribute (truncated to 200 chars).
+        attributes:     Static key-value attributes added to every span.
+
+    Returns:
+        Decorated callable (sync or async), or a single-argument decorator
+        when keyword arguments are supplied.
+    """
+
+    def _decorate(f: Callable[..., Any]) -> Callable[..., Any]:
+        effective_op = "execute_tool" if tool else operation
+        dec = _TraceDecorator(
+            fn=f,
+            name=name,
+            model=model,
+            operation=effective_op,
+            tool=tool,
+            capture_args=capture_args,
+            capture_return=capture_return,
+            attributes=attributes,
+        )
+        if inspect.iscoroutinefunction(f):
+            return _make_async_wrapper(dec, f)
+        return dec
+
+    if fn is not None:
+        # @trace — bare decorator, no parentheses
+        return _decorate(fn)
+    # @trace(...) — decorator factory
+    return _decorate