71 lines
1.6 KiB
Python
71 lines
1.6 KiB
Python
"""Unit tests configuration module."""
|
|
|
|
import os
|
|
|
|
import pytest
|
|
from openai import OpenAI
|
|
|
|
from opentelemetry import trace
|
|
from opentelemetry.instrumentation.openai_v2 import OpenAIInstrumentor
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
|
|
InMemorySpanExporter,
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def exporter():
|
|
exporter = InMemorySpanExporter()
|
|
processor = SimpleSpanProcessor(exporter)
|
|
|
|
provider = TracerProvider()
|
|
provider.add_span_processor(processor)
|
|
trace.set_tracer_provider(provider)
|
|
|
|
return exporter
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_exporter(exporter):
|
|
exporter.clear()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def environment():
|
|
if not os.getenv("OPENAI_API_KEY"):
|
|
os.environ["OPENAI_API_KEY"] = "test-api-key"
|
|
|
|
|
|
@pytest.fixture
|
|
def openai_client():
|
|
return OpenAI()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def vcr_config():
|
|
return {
|
|
"filter_headers": ["authorization", "api-key"],
|
|
"decode_compressed_response": True,
|
|
"before_record_response": scrub_response_headers,
|
|
}
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def instrument():
|
|
OpenAIInstrumentor().instrument()
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def uninstrument():
|
|
OpenAIInstrumentor().uninstrument()
|
|
|
|
|
|
def scrub_response_headers(response):
|
|
"""
|
|
This scrubs sensitive response headers. Note they are case-sensitive!
|
|
"""
|
|
response["headers"]["openai-organization"] = "test_organization"
|
|
response["headers"]["Set-Cookie"] = "test_set_cookie"
|
|
return response
|