tests/test_retriever.py

# tests/test_retriever.py
import pytest
from modules.retriever import Retriever

def test_retriever_init():
    retriever = Retriever(api_key="test_key")
    assert retriever.client.api_key == "test_key"

def test_search_returns_list(monkeypatch):
    class MockResponse:
        def search(self, *args, **kwargs):
            return {
                "results": [
                    {"title": "Test Result 1", "url": "https://example.com/1"},
                    {"title": "Test Result 2", "url": "https://example.com/2"},
                ]
            }

    def mock_tavily_init(*args, **kwargs):
        return MockResponse()

    monkeypatch.setattr("modules.retriever.TavilyClient", mock_tavily_init)

    retriever = Retriever(api_key="test_key")
    results = retriever.search("test query")
    assert len(results) == 2
    assert results[0]["title"] == "Test Result 1"

def test_get_related_queries():
    retriever = Retriever(api_key="test_key")
    queries = retriever.get_related_queries("AI research")
    assert len(queries) == 3
    assert "AI research research paper" in queries