Skip to content

Commit 0a17fa6

Browse files
authored
Add embedder factory tests (#7) (#14)
Covers Embedder.from_callable (shapes, dtype, embed_one) and from_endpoint in both single-request and batch modes against a stdlib mock HTTP server, including custom-header pass-through. The from_sentence_transformers test is guarded with pytest.importorskip so it skips when the optional sentence-transformers extra is not installed, keeping core CI fast.
1 parent b43722b commit 0a17fa6

1 file changed

Lines changed: 140 additions & 0 deletions

File tree

‎tests/test_embedder.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
"""Tests for the Embedder factories (embeddings/embedder.py).
2+
3+
`from_callable` and `from_endpoint` are exercised with no optional
4+
dependencies: the endpoint factory is tested against a stdlib mock HTTP
5+
server. The `from_sentence_transformers` factory pulls in the optional
6+
`sentence-transformers` extra, so its test is skipped when that extra is not
7+
installed, keeping the core CI matrix fast.
8+
"""
9+
10+
import json
11+
import threading
12+
from http.server import BaseHTTPRequestHandler, HTTPServer
13+
from pathlib import Path
14+
15+
import numpy as np
16+
import pytest
17+
import sys
18+
19+
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
20+
21+
from tracer.embeddings.embedder import Embedder
22+
23+
24+
# ── from_callable ─────────────────────────────────────────────────────────────
25+
26+
def test_from_callable_shapes_and_dtype():
27+
def fn(texts):
28+
return [[float(len(t)), 1.0] for t in texts]
29+
30+
emb = Embedder.from_callable(fn)
31+
out = emb.embed(["a", "bbb"])
32+
assert out.shape == (2, 2)
33+
assert out.dtype == np.float32
34+
np.testing.assert_allclose(out[:, 0], [1.0, 3.0])
35+
36+
37+
def test_embed_one_returns_1d():
38+
emb = Embedder.from_callable(lambda texts: [[1.0, 2.0, 3.0] for _ in texts])
39+
one = emb.embed_one("hello")
40+
assert one.shape == (3,)
41+
assert one.dtype == np.float32
42+
43+
44+
def test_from_callable_repr():
45+
emb = Embedder.from_callable(lambda texts: [[0.0] for _ in texts])
46+
assert "callable" in repr(emb)
47+
48+
49+
# ── from_endpoint (mock HTTP server) ───────────────────────────────────────────
50+
51+
class _EmbedHandler(BaseHTTPRequestHandler):
52+
"""Echoes a deterministic embedding derived from input length.
53+
54+
Single mode: {"input": "ab"} -> {"embedding": [2.0]}
55+
Batch mode: {"inputs": ["a", "bb"]} -> {"embeddings": [[1.0], [2.0]]}
56+
"""
57+
58+
def do_POST(self):
59+
length = int(self.headers.get("Content-Length", 0))
60+
body = json.loads(self.rfile.read(length))
61+
self.server.request_count += 1
62+
self.server.last_auth = self.headers.get("Authorization")
63+
64+
if "inputs" in body:
65+
payload = {"embeddings": [[float(len(t))] for t in body["inputs"]]}
66+
else:
67+
payload = {"embedding": [float(len(body["input"]))]}
68+
69+
raw = json.dumps(payload).encode()
70+
self.send_response(200)
71+
self.send_header("Content-Type", "application/json")
72+
self.send_header("Content-Length", str(len(raw)))
73+
self.end_headers()
74+
self.wfile.write(raw)
75+
76+
def log_message(self, *args):
77+
pass
78+
79+
80+
@pytest.fixture
81+
def mock_endpoint():
82+
httpd = HTTPServer(("127.0.0.1", 0), _EmbedHandler)
83+
httpd.request_count = 0
84+
httpd.last_auth = None
85+
thread = threading.Thread(target=httpd.serve_forever, daemon=True)
86+
thread.start()
87+
port = httpd.server_address[1]
88+
yield httpd, f"http://127.0.0.1:{port}"
89+
httpd.shutdown()
90+
thread.join(timeout=5)
91+
httpd.server_close()
92+
93+
94+
def test_from_endpoint_single_mode(mock_endpoint):
95+
httpd, url = mock_endpoint
96+
emb = Embedder.from_endpoint(url)
97+
out = emb.embed(["a", "bb", "ccc"])
98+
assert out.shape == (3, 1)
99+
assert out.dtype == np.float32
100+
np.testing.assert_allclose(out[:, 0], [1.0, 2.0, 3.0])
101+
# One HTTP request per text in single mode.
102+
assert httpd.request_count == 3
103+
104+
105+
def test_from_endpoint_batch_mode(mock_endpoint):
106+
httpd, url = mock_endpoint
107+
emb = Embedder.from_endpoint(url, batch_key="inputs",
108+
batch_output_key="embeddings")
109+
out = emb.embed(["a", "bb"])
110+
assert out.shape == (2, 1)
111+
np.testing.assert_allclose(out[:, 0], [1.0, 2.0])
112+
# The whole batch goes in a single request.
113+
assert httpd.request_count == 1
114+
115+
116+
def test_from_endpoint_passes_headers(mock_endpoint):
117+
httpd, url = mock_endpoint
118+
emb = Embedder.from_endpoint(url, headers={"Authorization": "Bearer xyz"})
119+
emb.embed(["a"])
120+
assert httpd.last_auth == "Bearer xyz"
121+
122+
123+
def test_from_endpoint_repr(mock_endpoint):
124+
_, url = mock_endpoint
125+
emb = Embedder.from_endpoint(url)
126+
assert "endpoint" in repr(emb)
127+
128+
129+
# ── from_sentence_transformers (optional extra) ────────────────────────────────
130+
131+
def test_from_sentence_transformers_smoke():
132+
# Skipped unless the optional `sentence-transformers` extra is installed,
133+
# so the core CI matrix stays fast.
134+
pytest.importorskip("sentence_transformers")
135+
emb = Embedder.from_sentence_transformers("all-MiniLM-L6-v2")
136+
out = emb.embed(["hello world", "another sentence"])
137+
assert out.ndim == 2
138+
assert out.shape[0] == 2
139+
assert out.dtype == np.float32
140+
assert "sentence-transformers" in repr(emb)

0 commit comments

Comments
 (0)