90 lines
No EOL
2.1 KiB
Python
90 lines
No EOL
2.1 KiB
Python
"""Pytest configuration and fixtures."""
|
|
|
|
import os
|
|
import sys
|
|
import warnings
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from dotenv import load_dotenv
|
|
|
|
# Suppress warnings
|
|
warnings.filterwarnings("ignore", category=DeprecationWarning, module="cohere")
|
|
warnings.filterwarnings("ignore", message=".*__fields__.*", category=DeprecationWarning)
|
|
|
|
# Add router to path
|
|
sys.path.insert(0, str(Path(__file__).parent))
|
|
|
|
# Load environment variables once
|
|
load_dotenv()
|
|
|
|
|
|
def pytest_addoption(parser):
|
|
"""Add custom command line options."""
|
|
parser.addoption(
|
|
"--ollama",
|
|
action="store_true",
|
|
default=False,
|
|
help="Run Ollama tests"
|
|
)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def api_keys():
|
|
"""Fixture to provide API keys."""
|
|
return {
|
|
"cohere": os.getenv("COHERE_API_KEY"),
|
|
"gemini": os.getenv("GEMINI_API_KEY"),
|
|
"openai": os.getenv("OPENAI_API_KEY"),
|
|
"openai_base_url": os.getenv("OPENAI_BASE_URL", "https://veronica.pratikn.com"),
|
|
"ollama_base_url": os.getenv("OLLAMA_BASE_URL", "http://localhost:11434"),
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def cohere_embedder():
|
|
"""Create a Cohere embedding instance."""
|
|
from router.embed import CohereEmbedding
|
|
return CohereEmbedding()
|
|
|
|
|
|
@pytest.fixture
|
|
def gemini_embedder():
|
|
"""Create a Gemini embedding instance."""
|
|
from router.embed import GeminiEmbedding
|
|
return GeminiEmbedding()
|
|
|
|
|
|
@pytest.fixture
|
|
def ollama_embedder():
|
|
"""Create an Ollama embedding instance."""
|
|
from router.embed import OllamaEmbedding
|
|
return OllamaEmbedding()
|
|
|
|
|
|
@pytest.fixture
|
|
def cohere_reranker():
|
|
"""Create a Cohere rerank instance."""
|
|
from router.rerank import CohereRerank
|
|
return CohereRerank()
|
|
|
|
|
|
@pytest.fixture
|
|
def openai_router(api_keys):
|
|
"""Create an OpenAI router instance."""
|
|
from router.openai_compatible import OpenAI
|
|
return OpenAI(base_url=api_keys["openai_base_url"])
|
|
|
|
|
|
@pytest.fixture
|
|
def gemini_router():
|
|
"""Create a Gemini router instance."""
|
|
from router.gemini import Gemini
|
|
return Gemini()
|
|
|
|
|
|
@pytest.fixture
|
|
def cohere_router():
|
|
"""Create a Cohere router instance."""
|
|
from router.cohere import Cohere
|
|
return Cohere() |