tests/unit/test_face_index.py

"""
Tests for FaceIndex — the reverse face search index.

Tests run without needing ONNX models or the full app — they use
synthetic embeddings to verify storage/retrieval logic.
"""
from __future__ import annotations

import sys
import os
import numpy as np
import pytest

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from storage.face_index import FaceIndex


@pytest.fixture
def index():
    """Fresh in-memory FaceIndex for each test."""
    return FaceIndex(":memory:")


def _random_embedding(seed: int) -> np.ndarray:
    """Generate a deterministic L2-normalized random embedding."""
    rng = np.random.RandomState(seed)
    emb = rng.randn(512).astype(np.float32)
    return emb / np.linalg.norm(emb)


def _similar_embedding(base: np.ndarray, noise_scale: float = 0.1, seed: int = 0) -> np.ndarray:
    """Generate an embedding similar to `base` with small noise."""
    rng = np.random.RandomState(seed)
    emb = base + noise_scale * rng.randn(512)
    return (emb / np.linalg.norm(emb)).astype(np.float32)


# --------------------------------------------------------------------------- #
# Initialization
# --------------------------------------------------------------------------- #

def test_index_initializes(index):
    assert index.count() == 0
    assert index.EMBEDDING_DIM == 512


def test_index_stats_empty(index):
    stats = index.stats()
    assert stats["total_faces"] == 0
    assert stats["named_faces"] == 0
    assert stats["embedding_dim"] == 512


# --------------------------------------------------------------------------- #
# Enrollment
# --------------------------------------------------------------------------- #

def test_enroll_returns_face_id(index):
    emb = _random_embedding(42)
    face_id = index.enroll(emb, name="Alice")
    assert isinstance(face_id, str)
    assert len(face_id) > 0


def test_enroll_increments_count(index):
    emb1 = _random_embedding(1)
    emb2 = _random_embedding(2)
    index.enroll(emb1, name="Alice")
    assert index.count() == 1
    index.enroll(emb2, name="Bob")
    assert index.count() == 2


def test_enroll_with_metadata(index):
    emb = _random_embedding(42)
    face_id = index.enroll(
        emb, name="Alice",
        source_url="https://example.com/alice.jpg",
        metadata={"age": 30, "location": "US"},
    )
    face = index.get(face_id)
    assert face is not None
    assert face["name"] == "Alice"
    assert face["source_url"] == "https://example.com/alice.jpg"
    assert face["metadata"]["age"] == 30
    assert face["metadata"]["location"] == "US"


def test_enroll_rejects_wrong_shape(index):
    bad_emb = np.zeros(256, dtype=np.float32)
    with pytest.raises(ValueError, match="shape"):
        index.enroll(bad_emb)


def test_enroll_l2_normalizes(index):
    """Even if the input isn't L2-normalized, the index should handle it."""
    emb = np.ones(512, dtype=np.float32) * 5  # not normalized
    face_id = index.enroll(emb, name="Test")
    assert face_id  # should not raise


# --------------------------------------------------------------------------- #
# Search
# --------------------------------------------------------------------------- #

def test_search_returns_best_match_first(index):
    """The most similar face should be returned first."""
    emb_alice = _random_embedding(1)
    emb_bob = _random_embedding(2)
    index.enroll(emb_alice, name="Alice")
    index.enroll(emb_bob, name="Bob")

    # Query is similar to Alice
    query = _similar_embedding(emb_alice, noise_scale=0.1)
    results = index.search(query, top_k=5, threshold=0.0)

    assert len(results) >= 1
    assert results[0]["name"] == "Alice"
    assert results[0]["similarity"] > results[1]["similarity"] if len(results) > 1 else True


def test_search_returns_empty_when_index_empty(index):
    query = _random_embedding(42)
    results = index.search(query, top_k=5, threshold=0.0)
    assert results == []


def test_search_threshold_filters_results(index):
    """Threshold should filter out low-similarity matches."""
    emb_alice = _random_embedding(1)
    emb_bob = _random_embedding(2)
    index.enroll(emb_alice, name="Alice")
    index.enroll(emb_bob, name="Bob")

    # Query is similar to Alice (sim ~0.95) and dissimilar to Bob (sim ~0.0)
    query = _similar_embedding(emb_alice, noise_scale=0.05)

    # Threshold 0.5 — only Alice should match
    results = index.search(query, top_k=5, threshold=0.5)
    assert all(r["similarity"] >= 0.5 for r in results)
    assert any(r["name"] == "Alice" for r in results)
    assert not any(r["name"] == "Bob" for r in results)


def test_search_top_k_limits_results(index):
    for i in range(10):
        emb = _random_embedding(i)
        index.enroll(emb, name=f"Person_{i}")

    query = _random_embedding(0)
    results = index.search(query, top_k=3, threshold=0.0)
    assert len(results) <= 3


def test_search_similarity_values_match_numpy(index):
    """Cosine similarity returned by sqlite-vec should match numpy's dot product."""
    emb1 = _random_embedding(1)
    emb2 = _random_embedding(2)
    index.enroll(emb1, name="Alice")
    index.enroll(emb2, name="Bob")

    query = _similar_embedding(emb1, noise_scale=0.1)

    # Direct numpy cosine sim
    expected_sim_alice = float(np.dot(query, emb1))

    results = index.search(query, top_k=5, threshold=0.0)
    alice_result = next(r for r in results if r["name"] == "Alice")
    assert abs(alice_result["similarity"] - expected_sim_alice) < 0.01


# --------------------------------------------------------------------------- #
# CRUD
# --------------------------------------------------------------------------- #

def test_get_returns_face(index):
    emb = _random_embedding(42)
    face_id = index.enroll(emb, name="Alice", source_url="https://example.com")
    face = index.get(face_id)
    assert face is not None
    assert face["name"] == "Alice"
    assert face["source_url"] == "https://example.com"


def test_get_returns_none_for_missing(index):
    assert index.get("nonexistent-id") is None


def test_list_returns_all(index):
    for i in range(5):
        index.enroll(_random_embedding(i), name=f"Person_{i}")
    faces = index.list(limit=10)
    assert len(faces) == 5


def test_list_filter_by_name(index):
    index.enroll(_random_embedding(1), name="Alice")
    index.enroll(_random_embedding(2), name="Bob")
    index.enroll(_random_embedding(3), name="Alice Cooper")

    faces = index.list(name="Alice")
    assert len(faces) == 2  # "Alice" and "Alice Cooper"


def test_delete_removes_face(index):
    face_id = index.enroll(_random_embedding(42), name="Alice")
    assert index.count() == 1
    deleted = index.delete(face_id)
    assert deleted is True
    assert index.count() == 0
    assert index.get(face_id) is None


def test_delete_returns_false_for_missing(index):
    assert index.delete("nonexistent-id") is False


def test_clear_removes_all(index):
    for i in range(5):
        index.enroll(_random_embedding(i), name=f"Person_{i}")
    assert index.count() == 5
    n = index.clear()
    assert n == 5
    assert index.count() == 0


# --------------------------------------------------------------------------- #
# Stats
# --------------------------------------------------------------------------- #

def test_stats_after_enrollment(index):
    index.enroll(_random_embedding(1), name="Alice")
    index.enroll(_random_embedding(2))  # anonymous

    stats = index.stats()
    assert stats["total_faces"] == 2
    assert stats["named_faces"] == 1
    assert stats["anonymous_faces"] == 1
    assert stats["last_enrollment"] is not None
    assert stats["recent_enrollments_24h"] == 2