ReviewLlama/tests/test_vector_store.py

import os
import tempfile
from pathlib import Path
from unittest.mock import Mock, patch

import pytest
from langchain_core.documents.base import Document
from langchain_core.vectorstores import VectorStoreRetriever

from reviewllama.vector_store import create_retriever, documents_from_path_list


@pytest.fixture
def temp_files():
    """Create temporary test files"""
    temp_dir = tempfile.mkdtemp()
    files = []

    # Create test files with different content
    file_contents = [
        "This is the first test file with Python code examples.",
        "The second file contains JavaScript functions and classes.",
        "File three has documentation about testing best practices.",
        "Final file includes configuration settings and deployment info.",
    ]

    for i, content in enumerate(file_contents):
        file_path = Path(temp_dir) / f"test_file_{i}.txt"
        file_path.write_text(content)
        files.append(file_path)

    yield files

    # Cleanup
    for file_path in files:
        file_path.unlink()
    os.rmdir(temp_dir)


def test_load_documents(temp_files):
    docs = documents_from_path_list(temp_files)

    assert len(docs) == 4
    assert all(isinstance(doc, Document) for doc in docs)
    assert all(doc.page_content for doc in docs)
    assert all(doc.metadata.get("source") for doc in docs)


@patch("reviewllama.vector_store.OllamaEmbeddings")
@patch("reviewllama.vector_store.FAISS")
@patch("reviewllama.vector_store.documents_from_path_list")
def test_create_retriever(
    mock_docs_from_list, mock_faiss, mock_embeddings, ollama_config
):
    """Test successful retriever creation"""
    # Setup mocks
    mock_docs = [Document(page_content="test", metadata={"source": "test.txt"})]
    mock_docs_from_list.return_value = mock_docs

    mock_embedding_instance = Mock()
    mock_embeddings.return_value = mock_embedding_instance

    mock_vectorstore = Mock()
    mock_retriever = Mock(spec=VectorStoreRetriever)
    mock_vectorstore.as_retriever.return_value = mock_retriever
    mock_faiss.from_documents.return_value = mock_vectorstore

    # Test
    result = create_retriever(["test.txt"], ollama_config)

    # Assertions
    assert result == mock_retriever
    mock_embeddings.assert_called_once_with(
        model=ollama_config.embedding_model, base_url=ollama_config.base_url
    )
    mock_docs_from_list.assert_called_once_with(["test.txt"])
    mock_faiss.from_documents.assert_called_once_with(
        mock_docs, mock_embedding_instance
    )
    mock_vectorstore.as_retriever.assert_called_once()