tests/test_index.py

"""Tests for the vector index."""

from __future__ import annotations

from pathlib import Path

import pytest

from mathvision_explorer.index import VectorIndex


def test_vector_index_searches_by_cosine_similarity() -> None:
    """Nearest-neighbor search ranks vectors by cosine similarity."""

    index = VectorIndex()
    index.add("same", (1.0, 0.0))
    index.add("opposite", (-1.0, 0.0))
    index.add("orthogonal", (0.0, 1.0))

    neighbors = index.search((1.0, 0.0), limit=2)

    assert [neighbor.item_id for neighbor in neighbors] == ["same", "orthogonal"]
    assert neighbors[0].score == pytest.approx(1.0)


def test_vector_index_round_trips_tsv(tmp_path: Path) -> None:
    """Saved indexes can be loaded back from disk."""

    path = tmp_path / "index.tsv"
    index = VectorIndex()
    index.add("a", (0.25, 0.75))

    index.save_tsv(path)
    loaded = VectorIndex.load_tsv(path)

    assert len(loaded) == 1
    assert loaded.search((0.25, 0.75))[0].item_id == "a"


def test_vector_dimensions_must_match() -> None:
    """Search vectors must use the same dimensionality as stored vectors."""

    index = VectorIndex()
    index.add("a", (1.0, 2.0))

    with pytest.raises(ValueError, match="same dimensions"):
        index.search((1.0,))