| """Unit tests for cores.embedding — vector ops + cache.""" |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from cores.embedding import ( |
| normalize, cosine_similarity, euclidean_distance, |
| batch_cosine_similarity, EmbeddingCache, |
| ) |
|
|
|
|
| class TestVectors: |
| def test_normalize(self): |
| v = np.array([3.0, 4.0]) |
| n = normalize(v) |
| assert np.linalg.norm(n) == pytest_approx(1.0) if (pytest := __import__("pytest")) else True |
|
|
| def test_normalize_zero_vector(self): |
| v = np.zeros(3) |
| n = normalize(v) |
| assert np.all(n == 0) |
|
|
| def test_cosine_similarity_identical(self): |
| v = np.array([1.0, 2.0, 3.0]) |
| assert cosine_similarity(v, v) == 1.0 |
|
|
| def test_euclidean_distance(self): |
| a = np.array([0.0, 0.0]) |
| b = np.array([3.0, 4.0]) |
| assert euclidean_distance(a, b) == 5.0 |
|
|
| def test_batch_cosine_similarity(self): |
| query = np.array([1.0, 0.0]) |
| matrix = np.array([ |
| [1.0, 0.0], |
| [0.0, 1.0], |
| [-1.0, 0.0], |
| ]) |
| sims = batch_cosine_similarity(query, matrix) |
| assert len(sims) == 3 |
| assert sims[0] == 1.0 |
| assert sims[1] == 0.0 |
| assert sims[2] == -1.0 |
|
|
|
|
| class TestEmbeddingCache: |
| def test_get_or_load_loads_once(self): |
| cache = EmbeddingCache() |
| call_count = [0] |
| def loader(): |
| call_count[0] += 1 |
| return {"model": "fake"} |
| m1 = cache.get_or_load("key1", loader) |
| m2 = cache.get_or_load("key1", loader) |
| assert m1 is m2 |
| assert call_count[0] == 1 |
|
|
| def test_is_loaded(self): |
| cache = EmbeddingCache() |
| assert not cache.is_loaded("x") |
| cache.get_or_load("x", lambda: "model") |
| assert cache.is_loaded("x") |
|
|
| def test_evict(self): |
| cache = EmbeddingCache() |
| cache.get_or_load("x", lambda: "model") |
| assert cache.evict("x") is True |
| assert not cache.is_loaded("x") |
| assert cache.evict("x") is False |
|
|
| def test_clear(self): |
| cache = EmbeddingCache() |
| cache.get_or_load("a", lambda: 1) |
| cache.get_or_load("b", lambda: 2) |
| n = cache.clear() |
| assert n == 2 |
| assert cache.keys() == [] |
|
|
| def test_keys(self): |
| cache = EmbeddingCache() |
| cache.get_or_load("a", lambda: 1) |
| cache.get_or_load("b", lambda: 2) |
| assert set(cache.keys()) == {"a", "b"} |
|
|
|
|
| def pytest_approx(expected, rel=1e-6): |
| """Tiny local approx since pytest.approx may not be in scope.""" |
| class _Approx: |
| def __init__(self, expected, rel): |
| self.expected = expected |
| self.rel = rel |
| def __eq__(self, other): |
| return abs(other - self.expected) <= self.rel * max(abs(self.expected), abs(other), 1.0) |
| return _Approx(expected, rel) |
|
|
