File size: 2,885 Bytes
892fa81 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """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], # identical
[0.0, 1.0], # orthogonal
[-1.0, 0.0], # opposite
])
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)
|