tests/test_retrieval.py

"""
Tests for FAISS index and retrieval engine.

These tests verify:
- FAISS index building
- Search result shapes
- Similarity score validity
- Timing measurement
- Save/load roundtrip
"""

import pytest
import torch
import numpy as np
import tempfile
from pathlib import Path

from src.retrieval.index import FAISSIndex
from src.retrieval.engine import RetrievalEngine, RetrievalResult


@pytest.fixture
def dummy_embeddings():
    """Create L2-normalized dummy embeddings."""
    n_gallery = 100
    embed_dim = 768
    
    embeddings = torch.randn(n_gallery, embed_dim)
    return torch.nn.functional.normalize(embeddings, dim=1)


@pytest.fixture
def dummy_query():
    """Create a dummy query embedding."""
    query = torch.randn(768)
    return torch.nn.functional.normalize(query, dim=0)


@pytest.fixture
def built_index(dummy_embeddings):
    """Create a FAISSIndex with embeddings loaded."""
    index = FAISSIndex(embed_dim=768)
    index.build(dummy_embeddings)
    return index


class TestFAISSIndex:
    """Tests for FAISSIndex class."""
    
    def test_initialization(self):
        """Test index initializes correctly."""
        index = FAISSIndex(embed_dim=512)
        assert index.embed_dim == 512
        assert index.size == 0
        assert not index.is_built
    
    def test_build(self, dummy_embeddings):
        """Test index builds from embeddings."""
        index = FAISSIndex(embed_dim=768)
        index.build(dummy_embeddings)
        
        assert index.is_built
        assert index.size == 100
    
    def test_search_returns_correct_shape(self, built_index, dummy_query):
        """Test search returns correct shapes."""
        k = 5
        scores, indices = built_index.search(dummy_query, k=k)
        
        assert scores.shape == (1, k)
        assert indices.shape == (1, k)
    
    def test_search_scores_valid(self, built_index, dummy_query):
        """Test similarity scores are in valid range."""
        scores, _ = built_index.search(dummy_query, k=5)
        
        # Cosine similarity should be between -1 and 1
        # For L2-normalized vectors, typically between 0 and 1
        assert scores.min() >= -1.0
        assert scores.max() <= 1.0
    
    def test_search_indices_valid(self, built_index, dummy_query):
        """Test returned indices are valid."""
        _, indices = built_index.search(dummy_query, k=5)
        
        # All indices should be valid
        assert (indices >= 0).all()
        assert (indices < built_index.size).all()
    
    def test_search_k_clamped(self, built_index):
        """Test k is clamped to index size."""
        query = torch.randn(768)
        query = torch.nn.functional.normalize(query, dim=0)
        
        # Try to get more results than available
        scores, indices = built_index.search(query, k=200)
        
        # Should return at most index.size results
        assert scores.shape[1] <= built_index.size
    
    def test_search_before_build_raises(self):
        """Test search before build raises error."""
        index = FAISSIndex(embed_dim=768)
        query = torch.randn(768)
        
        with pytest.raises(RuntimeError):
            index.search(query, k=5)
    
    def test_save_load_roundtrip(self, dummy_embeddings):
        """Test save then load returns same results."""
        index = FAISSIndex(embed_dim=768)
        index.build(dummy_embeddings)
        
        query = torch.randn(768)
        query = torch.nn.functional.normalize(query, dim=0)
        
        # Get original results
        scores1, indices1 = index.search(query, k=5)
        
        with tempfile.TemporaryDirectory() as tmpdir:
            save_path = Path(tmpdir) / "test_index.faiss"
            
            # Save
            index.save(save_path)
            assert save_path.exists()
            
            # Load
            loaded_index = FAISSIndex(embed_dim=768)
            loaded_index.load(save_path)
            
            assert loaded_index.size == index.size
            
            # Search should return same results
            scores2, indices2 = loaded_index.search(query, k=5)
            
            assert np.allclose(scores1, scores2)
            assert np.array_equal(indices1, indices2)
    
    def test_get_embeddings(self, built_index):
        """Test get_embeddings returns correct shape."""
        embeddings = built_index.get_embeddings()
        
        assert embeddings.shape == (100, 768)
    
    def test_get_embeddings_normalized(self, built_index):
        """Test retrieved embeddings are L2-normalized."""
        embeddings = built_index.get_embeddings()
        
        norms = np.linalg.norm(embeddings, axis=1)
        assert np.allclose(norms, 1.0, atol=1e-5)


class TestRetrievalEngine:
    """Tests for RetrievalEngine class."""
    
    def test_timing_stats_empty(self):
        """Test timing stats with no queries."""
        engine = RetrievalEngine.__new__(RetrievalEngine)
        engine._query_times_list = []
        
        stats = engine.get_timing_stats()
        
        assert stats["mean"] == 0
        assert stats["median"] == 0
    
    def test_timing_stats_computed(self):
        """Test timing stats are computed correctly."""
        engine = RetrievalEngine.__new__(RetrievalEngine)
        engine._query_times_list = [10.0, 20.0, 30.0, 40.0, 50.0]
        
        stats = engine.get_timing_stats()
        
        assert stats["mean"] == 30.0
        assert stats["median"] == 30.0
    
    def test_retrieval_result_structure(self):
        """Test RetrievalResult dataclass."""
        result = RetrievalResult(
            indices=[0, 1, 2],
            scores=[0.9, 0.8, 0.7],
            query_time_ms=15.5,
            modality="optical"
        )
        
        assert result.indices == [0, 1, 2]
        assert result.scores == [0.9, 0.8, 0.7]
        assert result.query_time_ms == 15.5
        assert result.modality == "optical"


class TestEndToEnd:
    """End-to-end tests with dummy data."""
    
    def test_index_search_pipeline(self, dummy_embeddings, dummy_query):
        """Test complete build-search pipeline."""
        # Build
        index = FAISSIndex(embed_dim=768)
        index.build(dummy_embeddings)
        
        # Search
        scores, indices = index.search(dummy_query, k=10)
        
        # Verify
        assert scores.shape == (1, 10)
        assert indices.shape == (1, 10)
        
        # Results should be sorted by score (descending)
        assert scores[0, 0] >= scores[0, -1]


# Self-check
if __name__ == "__main__":
    pytest.main([__file__, "-v"])