37 lines
1.1 KiB
Python
37 lines
1.1 KiB
Python
"""Pytest configuration for benchmark tests."""
|
|
|
|
import os
|
|
from collections.abc import Generator
|
|
|
|
import pytest
|
|
from langsmith import Client, get_tracing_context
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def langsmith_client() -> Generator[Client | None, None, None]:
|
|
"""Create a LangSmith client if LANGSMITH_API_KEY is set.
|
|
|
|
This fixture is session-scoped and automatically used by all tests.
|
|
It creates a single client instance and ensures it's flushed after each test.
|
|
"""
|
|
langsmith_api_key = os.environ.get("LANGSMITH_API_KEY") or os.environ.get("LANGCHAIN_API_KEY")
|
|
|
|
if langsmith_api_key:
|
|
client = get_tracing_context()["client"] or Client()
|
|
yield client
|
|
|
|
# Final flush at end of session
|
|
client.flush()
|
|
else:
|
|
yield None
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def flush_langsmith_after_test(langsmith_client: Client) -> Generator[None, None, None]:
|
|
"""Automatically flush LangSmith client after each test."""
|
|
yield
|
|
|
|
# This runs after each test
|
|
if langsmith_client is not None:
|
|
langsmith_client.flush()
|