|
| 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