tests/test_data_loader.py

"""
Unit tests for WikiData loader.

Note: These tests require the data files to be present. Tests will be
skipped if data files are missing.
"""

import numpy as np
import pytest

from src.config import EMBEDDING_DIM, validate_data_files

# Skip all tests if data files are missing
pytestmark = pytest.mark.skipif(
    not all(validate_data_files().values()),
    reason="Data files not available",
)


@pytest.fixture(scope="module")
def wiki_data():
    """Load wiki_data once for all tests in this module."""
    from src.data import wiki_data

    # Trigger loading
    _ = wiki_data.article_count()
    return wiki_data


class TestCoreAccessors:
    """Test core accessor methods."""

    def test_article_count_positive(self, wiki_data):
        """Article count should be positive."""
        assert wiki_data.article_count() > 0

    def test_traversable_count_positive(self, wiki_data):
        """Traversable count should be positive."""
        assert wiki_data.traversable_count() > 0

    def test_traversable_less_than_total(self, wiki_data):
        """Traversable count should be less than total (not all have links)."""
        assert wiki_data.traversable_count() <= wiki_data.article_count()

    def test_get_title_valid_index(self, wiki_data):
        """Should return title for valid index."""
        title = wiki_data.get_title(0)
        assert isinstance(title, str)

    def test_get_title_invalid_index_raises(self, wiki_data):
        """Should raise IndexError for invalid index."""
        with pytest.raises(IndexError):
            wiki_data.get_title(-1)
        with pytest.raises(IndexError):
            wiki_data.get_title(999_999_999)

    def test_get_index_known_article(self, wiki_data):
        """Should return index for known article."""
        idx = wiki_data.get_index("Albert Einstein")
        assert idx is not None
        assert isinstance(idx, int)
        assert idx >= 0

    def test_get_index_unknown_article(self, wiki_data):
        """Should return None for unknown article."""
        idx = wiki_data.get_index("ZZZZZ_NOT_A_REAL_ARTICLE_12345")
        assert idx is None

    def test_has_article_true(self, wiki_data):
        """Should return True for existing article."""
        assert wiki_data.has_article("Albert Einstein") is True

    def test_has_article_false(self, wiki_data):
        """Should return False for non-existing article."""
        assert wiki_data.has_article("ZZZZZ_NOT_A_REAL_ARTICLE_12345") is False

    def test_is_traversable_true(self, wiki_data):
        """Should return True for article with outgoing links."""
        # Albert Einstein should have links
        assert wiki_data.is_traversable("Albert Einstein") is True

    def test_is_traversable_false_nonexistent(self, wiki_data):
        """Should return False for non-existent article."""
        assert wiki_data.is_traversable("ZZZZZ_NOT_A_REAL_ARTICLE_12345") is False


class TestEmbeddingAccessors:
    """Test embedding accessor methods."""

    def test_get_embedding_valid(self, wiki_data):
        """Should return embedding for valid article."""
        emb = wiki_data.get_embedding("Albert Einstein")
        assert emb is not None
        assert isinstance(emb, np.ndarray)
        assert emb.shape == (EMBEDDING_DIM,)

    def test_get_embedding_invalid(self, wiki_data):
        """Should return None for invalid article."""
        emb = wiki_data.get_embedding("ZZZZZ_NOT_A_REAL_ARTICLE_12345")
        assert emb is None

    def test_get_embedding_by_idx(self, wiki_data):
        """Should return embedding by index."""
        idx = wiki_data.get_index("Albert Einstein")
        emb = wiki_data.get_embedding_by_idx(idx)
        assert isinstance(emb, np.ndarray)
        assert emb.shape == (EMBEDDING_DIM,)

    def test_get_normalized_embedding(self, wiki_data):
        """Should return normalized embedding (L2 norm ~= 1)."""
        emb = wiki_data.get_normalized_embedding("Albert Einstein")
        assert emb is not None
        norm = np.linalg.norm(emb)
        assert 0.99 < norm < 1.01  # Should be approximately 1

    def test_get_embeddings_batch(self, wiki_data):
        """Should return batch of embeddings."""
        titles = ["Albert Einstein", "Physics", "Mathematics"]
        batch = wiki_data.get_embeddings_batch(titles)
        assert isinstance(batch, np.ndarray)
        assert batch.shape == (3, EMBEDDING_DIM)

    def test_get_embeddings_batch_partial(self, wiki_data):
        """Should handle batch with some invalid titles."""
        titles = ["Albert Einstein", "INVALID_ARTICLE", "Physics"]
        batch = wiki_data.get_embeddings_batch(titles)
        assert batch.shape[0] == 2  # Only 2 valid


class TestSimilarityFunctions:
    """Test similarity computation methods."""

    def test_similarity_valid(self, wiki_data):
        """Should return similarity for valid articles."""
        sim = wiki_data.similarity("Physics", "Mathematics")
        assert sim is not None
        assert isinstance(sim, float)
        assert -1 <= sim <= 1  # Cosine similarity range

    def test_similarity_invalid(self, wiki_data):
        """Should return None if either article is invalid."""
        assert wiki_data.similarity("Physics", "INVALID") is None
        assert wiki_data.similarity("INVALID", "Physics") is None

    def test_similarity_self(self, wiki_data):
        """Similarity of article with itself should be ~1."""
        sim = wiki_data.similarity("Albert Einstein", "Albert Einstein")
        assert sim is not None
        assert sim > 0.99

    def test_similarity_semantic_ordering(self, wiki_data):
        """Related articles should be more similar than unrelated."""
        sim_related = wiki_data.similarity("Physics", "Mathematics")
        sim_unrelated = wiki_data.similarity("Physics", "Pizza")
        assert sim_related is not None
        assert sim_unrelated is not None
        assert sim_related > sim_unrelated

    def test_rank_by_similarity(self, wiki_data):
        """Should rank candidates by similarity."""
        candidates = ["Mathematics", "Pizza", "Chemistry"]
        ranked = wiki_data.rank_by_similarity(candidates, "Physics")
        assert len(ranked) == 3
        assert all(isinstance(r, tuple) and len(r) == 2 for r in ranked)
        # Check sorted descending
        sims = [r[1] for r in ranked]
        assert sims == sorted(sims, reverse=True)

    def test_rank_by_similarity_empty(self, wiki_data):
        """Should return empty list for invalid target."""
        ranked = wiki_data.rank_by_similarity(["Physics"], "INVALID")
        assert ranked == []

    def test_nearest_neighbors(self, wiki_data):
        """Should return nearest neighbors."""
        neighbors = wiki_data.nearest_neighbors("Physics", k=5)
        assert len(neighbors) == 5
        assert all(isinstance(n, tuple) and len(n) == 2 for n in neighbors)
        # Should not include self
        assert all(n[0] != "Physics" for n in neighbors)

    def test_nearest_neighbors_invalid(self, wiki_data):
        """Should return empty list for invalid article."""
        neighbors = wiki_data.nearest_neighbors("INVALID", k=5)
        assert neighbors == []


class TestGraphAccessors:
    """Test graph accessor methods."""

    def test_get_links_valid(self, wiki_data):
        """Should return links for valid traversable article."""
        links = wiki_data.get_links("Albert Einstein")
        assert isinstance(links, list)
        assert len(links) > 0
        assert all(isinstance(link, str) for link in links)

    def test_get_links_invalid(self, wiki_data):
        """Should return empty list for invalid article."""
        links = wiki_data.get_links("INVALID_ARTICLE_12345")
        assert links == []

    def test_get_links_by_idx(self, wiki_data):
        """Should return link indices."""
        idx = wiki_data.get_index("Albert Einstein")
        if wiki_data.is_traversable("Albert Einstein"):
            links = wiki_data.get_links_by_idx(idx)
            assert isinstance(links, list)
            assert all(isinstance(i, int) for i in links)

    def test_get_popularity(self, wiki_data):
        """Should return popularity (inbound link count)."""
        # Popular articles should have high popularity
        pop = wiki_data.get_popularity("United States")
        assert isinstance(pop, int)
        assert pop >= 0
        # United States should be linked to by many articles
        # Note: This test may be slow first time due to on-demand calculation

    def test_get_popularity_invalid(self, wiki_data):
        """Should return 0 for invalid article."""
        pop = wiki_data.get_popularity("INVALID_ARTICLE_12345")
        assert pop == 0


class TestEdgeCases:
    """Test edge cases and error handling."""

    def test_empty_string_title(self, wiki_data):
        """Should handle empty string gracefully."""
        assert wiki_data.get_index("") is None or wiki_data.get_index("") is not None
        # Just shouldn't crash

    def test_unicode_title(self, wiki_data):
        """Should handle unicode titles."""
        # Try some unicode - may or may not exist
        wiki_data.get_index("日本")  # Japan in Japanese
        wiki_data.get_index("Москва")  # Moscow in Russian
        # Just shouldn't crash

    def test_special_characters(self, wiki_data):
        """Should handle special characters in titles."""
        wiki_data.get_index("C++")
        wiki_data.get_index("AC/DC")
        # Just shouldn't crash


class TestValidation:
    """Test validation and stats methods."""

    def test_validate_all_pass(self, wiki_data):
        """All validation checks should pass."""
        validation = wiki_data.validate()
        assert all(validation.values()), f"Failed checks: {validation}"

    def test_stats_complete(self, wiki_data):
        """Stats should include all expected keys."""
        stats = wiki_data.stats()
        expected_keys = [
            "total_articles",
            "traversable_articles",
            "embedding_dim",
            "embedding_dtype",
            "faiss_vectors",
        ]
        for key in expected_keys:
            assert key in stats, f"Missing key: {key}"