lhl/code-review-graph
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
Files
ecfd03a21ced51449ac31815d3bfe8f8df2b3bcc
code-review-graph/code-review-graph-main/tests/test_embeddings.py
T

1647 lines
64 KiB
Python
Raw Blame History

"""Tests for the embeddings module."""
import json
import os
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from unittest.mock import MagicMock, patch
import pytest
from code_review_graph.embeddings import (
LOCAL_DEFAULT_MODEL,
EmbeddingStore,
GoogleEmbeddingProvider,
LocalEmbeddingProvider,
MiniMaxEmbeddingProvider,
OpenAIEmbeddingProvider,
VoyageEmbeddingProvider,
_cosine_similarity,
_decode_vector,
_encode_vector,
_is_localhost_url,
_node_to_text,
get_provider,
)
from code_review_graph.graph import GraphNode
class TestVectorEncoding:
def test_roundtrip(self):
original = [1.0, 2.5, -3.14, 0.0, 100.0]
blob = _encode_vector(original)
decoded = _decode_vector(blob)
assert len(decoded) == len(original)
for a, b in zip(original, decoded):
assert abs(a - b) < 1e-5
def test_empty_vector(self):
blob = _encode_vector([])
decoded = _decode_vector(blob)
assert decoded == []
def test_blob_size(self):
vec = [1.0, 2.0, 3.0]
blob = _encode_vector(vec)
assert len(blob) == 12 # 3 floats * 4 bytes each
class TestCosineSimilarity:
def test_identical_vectors(self):
v = [1.0, 2.0, 3.0]
assert abs(_cosine_similarity(v, v) - 1.0) < 1e-6
def test_orthogonal_vectors(self):
a = [1.0, 0.0]
b = [0.0, 1.0]
assert abs(_cosine_similarity(a, b)) < 1e-6
def test_opposite_vectors(self):
a = [1.0, 0.0]
b = [-1.0, 0.0]
assert abs(_cosine_similarity(a, b) - (-1.0)) < 1e-6
def test_zero_vector(self):
a = [0.0, 0.0]
b = [1.0, 2.0]
assert _cosine_similarity(a, b) == 0.0
def test_dimension_mismatch(self):
a = [1.0, 2.0, 3.0]
b = [1.0, 2.0]
assert _cosine_similarity(a, b) == 0.0
class TestNodeToText:
def _make_node(self, **kwargs):
defaults = dict(
id=1, kind="Function", name="my_func",
qualified_name="file.py::my_func", file_path="file.py",
line_start=1, line_end=10, language="python",
parent_name=None, params=None, return_type=None,
is_test=False, file_hash=None, extra={},
)
defaults.update(kwargs)
return GraphNode(**defaults)
def test_basic_function(self):
node = self._make_node()
text = _node_to_text(node)
assert "my_func" in text
assert "function" in text
assert "python" in text
def test_method_with_parent(self):
node = self._make_node(parent_name="MyClass")
text = _node_to_text(node)
assert "in MyClass" in text
def test_with_params_and_return_type(self):
node = self._make_node(params="(x: int, y: str)", return_type="bool")
text = _node_to_text(node)
assert "(x: int, y: str)" in text
assert "returns bool" in text
def test_file_node_no_kind(self):
node = self._make_node(kind="File", name="file.py")
text = _node_to_text(node)
# File kind should not add "file" as a kind label
assert "file.py" in text
class TestEmbeddingStore:
def _make_node(self, index: int) -> GraphNode:
return GraphNode(
id=index,
kind="Function",
name=f"func_{index}",
qualified_name=f"file.py::func_{index}",
file_path="file.py",
line_start=index,
line_end=index + 1,
language="python",
parent_name=None,
params=None,
return_type=None,
is_test=False,
file_hash=None,
extra={},
)
def test_store_initializes(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None):
store = EmbeddingStore(db)
assert store.count() == 0
store.close()
def test_count_empty(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None):
store = EmbeddingStore(db)
assert store.count() == 0
store.close()
def test_embed_nodes_returns_zero_when_unavailable(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None):
store = EmbeddingStore(db)
result = store.embed_nodes([])
assert result == 0
store.close()
def test_search_returns_empty_when_unavailable(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None):
store = EmbeddingStore(db)
results = store.search("query")
assert results == []
store.close()
def test_remove_node(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None):
store = EmbeddingStore(db)
# Should not raise even if node doesn't exist
store.remove_node("nonexistent::func")
store.close()
def test_embed_nodes_commits_each_batch(self, tmp_path):
db = tmp_path / "embeddings.db"
class Provider:
name = "test:provider"
def __init__(self):
self.calls = []
def embed(self, texts):
self.calls.append(list(texts))
return [[float(len(self.calls)), 0.0] for _ in texts]
def embed_query(self, text):
return [0.0, 0.0]
@property
def dimension(self):
return 2
provider = Provider()
nodes = [self._make_node(i) for i in range(5)]
with patch("code_review_graph.embeddings.get_provider", return_value=provider):
store = EmbeddingStore(db)
embedded = store.embed_nodes(nodes, batch_size=2)
assert embedded == 5
assert [len(call) for call in provider.calls] == [2, 2, 1]
assert store.count() == 5
store.close()
def test_embed_nodes_preserves_completed_batches_on_later_failure(self, tmp_path):
db = tmp_path / "embeddings.db"
class Provider:
name = "test:provider"
def __init__(self):
self.calls = 0
def embed(self, texts):
self.calls += 1
if self.calls == 2:
raise RuntimeError("rate limited")
return [[1.0, 0.0] for _ in texts]
def embed_query(self, text):
return [0.0, 0.0]
@property
def dimension(self):
return 2
provider = Provider()
nodes = [self._make_node(i) for i in range(5)]
with patch("code_review_graph.embeddings.get_provider", return_value=provider):
store = EmbeddingStore(db)
with pytest.raises(RuntimeError, match="rate limited"):
store.embed_nodes(nodes, batch_size=2)
assert store.count() == 2
store.close()
class TestLocalEmbeddingProviderModelName:
"""Tests for configurable model name on LocalEmbeddingProvider."""
def test_dimension_uses_current_sentence_transformers_api(self):
class CurrentModel:
def get_embedding_dimension(self):
return 384
provider = LocalEmbeddingProvider(model_name="offline/current")
provider._model = CurrentModel()
assert provider.dimension == 384
def test_dimension_supports_sentence_transformers_3(self):
class LegacyModel:
def get_sentence_embedding_dimension(self):
return 768
provider = LocalEmbeddingProvider(model_name="offline/legacy")
provider._model = LegacyModel()
assert provider.dimension == 768
def test_default_model_name(self):
with patch.dict(os.environ, {}, clear=False):
os.environ.pop("CRG_EMBEDDING_MODEL", None)
provider = LocalEmbeddingProvider()
assert provider._model_name == LOCAL_DEFAULT_MODEL
assert provider.name == f"local:{LOCAL_DEFAULT_MODEL}"
def test_explicit_model_name(self):
with patch.dict(os.environ, {"CRG_EMBEDDING_MODEL": "should-be-ignored"}):
provider = LocalEmbeddingProvider(model_name="custom/model")
assert provider._model_name == "custom/model"
assert provider.name == "local:custom/model"
def test_env_var_fallback(self):
with patch.dict(os.environ, {"CRG_EMBEDDING_MODEL": "BAAI/bge-small-en-v1.5"}):
provider = LocalEmbeddingProvider()
assert provider._model_name == "BAAI/bge-small-en-v1.5"
assert provider.name == "local:BAAI/bge-small-en-v1.5"
class TestGoogleEmbeddingProviderRetryLogging:
def test_retryable_error_logs_attempt_fraction_then_succeeds(self, caplog):
fn = MagicMock(side_effect=[RuntimeError("429 rate limited"), [1.0]])
with patch("code_review_graph.embeddings.time.sleep") as sleep:
result = GoogleEmbeddingProvider._call_with_retry(fn)
assert result == [1.0]
assert fn.call_count == 2
sleep.assert_called_once_with(1)
assert "Gemini API retry 1/3 in 1s (RuntimeError)" in caplog.text
def test_non_retryable_error_logs_once_without_sleeping(self, caplog):
caplog.set_level("DEBUG")
fn = MagicMock(side_effect=ValueError("400 invalid request"))
with (
patch("code_review_graph.embeddings.time.sleep") as sleep,
pytest.raises(ValueError, match="400 invalid request"),
):
GoogleEmbeddingProvider._call_with_retry(fn)
fn.assert_called_once_with()
sleep.assert_not_called()
assert "Non-retryable Gemini API error: ValueError" in caplog.text
def test_exhausted_retries_log_each_attempt_and_final_error(self, caplog):
fn = MagicMock(side_effect=RuntimeError("503 unavailable"))
with (
patch("code_review_graph.embeddings.time.sleep") as sleep,
pytest.raises(RuntimeError, match="503 unavailable"),
):
GoogleEmbeddingProvider._call_with_retry(fn)
assert fn.call_count == 3
assert [call.args for call in sleep.call_args_list] == [(1,), (2,)]
assert "Gemini API retry 1/3 in 1s (RuntimeError)" in caplog.text
assert "Gemini API retry 2/3 in 2s (RuntimeError)" in caplog.text
assert "Gemini API request failed after 3 requests" in caplog.text
class TestGetProviderValidation:
"""Unknown provider names must raise instead of silently using local."""
@pytest.mark.parametrize("name", ["opnai", "cohere", "moonbase", "MoOnBase"])
def test_unknown_provider_raises(self, name):
with pytest.raises(ValueError, match="Unknown embedding provider"):
get_provider(name)
def test_unknown_provider_message_lists_valid_names(self):
with pytest.raises(ValueError) as exc_info:
get_provider("moonbase")
msg = str(exc_info.value)
assert "moonbase" in msg
assert "Valid: local, openai, google, minimax, voyage" in msg
def test_case_and_whitespace_normalized_for_openai(self):
"""' OPENAI ' must route to the openai branch (and fail on its
missing env vars), not fall through to the local default."""
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="Missing required environment"):
get_provider(" OPENAI ")
def test_case_normalized_for_minimax(self):
with patch.dict(os.environ, {
"MINIMAX_API_KEY": "fake",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}, clear=False):
with patch(
"code_review_graph.embeddings.MiniMaxEmbeddingProvider",
) as mock_cls:
mock_cls.return_value = MagicMock()
provider = get_provider("MiniMax")
assert provider is mock_cls.return_value
@patch("code_review_graph.embeddings.LocalEmbeddingProvider")
@patch("code_review_graph.embeddings._check_available", return_value=True)
def test_local_case_and_whitespace_normalized(self, _mock_available, mock_cls):
mock_cls.return_value = MagicMock()
assert get_provider(" Local ") is mock_cls.return_value
@patch("code_review_graph.embeddings.LocalEmbeddingProvider")
@patch("code_review_graph.embeddings._check_available", return_value=True)
def test_none_and_empty_default_to_local(self, _mock_available, mock_cls):
mock_cls.return_value = MagicMock()
assert get_provider(None) is mock_cls.return_value
assert get_provider("") is mock_cls.return_value
assert get_provider(" ") is mock_cls.return_value
@patch("code_review_graph.embeddings.OpenAIEmbeddingProvider")
def test_none_defaults_to_openai_when_env_configured(self, mock_cls):
"""When provider is omitted but OpenAI env vars are set, default to
the OpenAI-compatible provider instead of local (#551)."""
mock_cls.return_value = MagicMock()
env = {
"CRG_OPENAI_API_KEY": "fake-key",
"CRG_OPENAI_BASE_URL": "http://localhost:11434/v1",
"CRG_OPENAI_MODEL": "nomic-embed-text",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}
with patch.dict(os.environ, env, clear=False):
provider = get_provider(None)
assert provider is mock_cls.return_value
mock_cls.assert_called_once_with(
api_key="fake-key",
base_url="http://localhost:11434/v1",
model="nomic-embed-text",
dimension=None,
batch_size=None,
)
@patch("code_review_graph.embeddings.OpenAIEmbeddingProvider")
@patch("code_review_graph.embeddings.LocalEmbeddingProvider")
@patch("code_review_graph.embeddings._check_available", return_value=True)
def test_explicit_blank_stays_local_when_openai_env_configured(
self, _mock_available, local_cls, openai_cls,
):
"""Only an omitted provider should use the configured OpenAI default."""
local_cls.return_value = MagicMock()
env = {
"CRG_OPENAI_API_KEY": "fake-key",
"CRG_OPENAI_BASE_URL": "http://localhost:11434/v1",
"CRG_OPENAI_MODEL": "nomic-embed-text",
}
with patch.dict(os.environ, env, clear=True):
assert get_provider("") is local_cls.return_value
assert get_provider(" ") is local_cls.return_value
openai_cls.assert_not_called()
class TestGetProviderModel:
"""Tests for model parameter in get_provider()."""
@patch("code_review_graph.embeddings.LocalEmbeddingProvider")
@patch("code_review_graph.embeddings._check_available", return_value=True)
def test_local_passes_model(self, _mock_available, mock_cls):
mock_cls.return_value = MagicMock()
get_provider(provider=None, model="custom/model")
mock_cls.assert_called_once_with(model_name="custom/model")
@patch("code_review_graph.embeddings.LocalEmbeddingProvider")
@patch("code_review_graph.embeddings._check_available", return_value=True)
def test_local_default_passes_none(self, _mock_available, mock_cls):
mock_cls.return_value = MagicMock()
get_provider(provider=None, model=None)
mock_cls.assert_called_once_with(model_name=None)
@patch("code_review_graph.embeddings._check_available", return_value=False)
def test_local_unavailable_returns_none(self, _mock_available):
assert get_provider("local") is None
@patch("code_review_graph.embeddings._check_available", return_value=False)
def test_embedding_store_unavailable_without_local_dependency(
self, _mock_available, tmp_path,
):
db = tmp_path / "embeddings.db"
store = EmbeddingStore(db, provider="local")
try:
assert store.available is False
finally:
store.close()
class TestCloudProviderWarning:
"""Tests for the stderr warning before cloud provider use (#174)."""
def test_minimax_triggers_stderr_warning(self, capsys):
"""Using the MiniMax provider should print a warning to stderr
unless CRG_ACCEPT_CLOUD_EMBEDDINGS=1 is set."""
with patch.dict(os.environ, {"MINIMAX_API_KEY": "fake"}, clear=False):
os.environ.pop("CRG_ACCEPT_CLOUD_EMBEDDINGS", None)
with patch(
"code_review_graph.embeddings.MiniMaxEmbeddingProvider",
) as mock_cls:
mock_cls.return_value = MagicMock()
get_provider(provider="minimax")
captured = capsys.readouterr()
assert "minimax" in captured.err.lower()
assert "cloud" in captured.err.lower()
assert "sent to an external API" in captured.err
# Should NOT have written to stdout (would corrupt MCP stdio).
assert captured.out == ""
def test_google_triggers_stderr_warning(self, capsys):
with patch.dict(os.environ, {"GOOGLE_API_KEY": "fake"}, clear=False):
os.environ.pop("CRG_ACCEPT_CLOUD_EMBEDDINGS", None)
with patch(
"code_review_graph.embeddings.GoogleEmbeddingProvider",
) as mock_cls:
mock_cls.return_value = MagicMock()
get_provider(provider="google")
captured = capsys.readouterr()
assert "google" in captured.err.lower()
assert captured.out == ""
def test_accept_env_var_suppresses_warning(self, capsys):
"""Setting CRG_ACCEPT_CLOUD_EMBEDDINGS=1 silences the warning."""
with patch.dict(os.environ, {
"MINIMAX_API_KEY": "fake",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}, clear=False):
with patch(
"code_review_graph.embeddings.MiniMaxEmbeddingProvider",
) as mock_cls:
mock_cls.return_value = MagicMock()
get_provider(provider="minimax")
captured = capsys.readouterr()
assert captured.err == ""
assert captured.out == ""
def test_local_provider_never_warns(self, capsys):
"""Local (offline) provider must not trigger the cloud warning."""
with patch(
"code_review_graph.embeddings.LocalEmbeddingProvider",
) as mock_cls:
with patch("code_review_graph.embeddings._check_available", return_value=True):
mock_cls.return_value = MagicMock()
get_provider(provider=None)
captured = capsys.readouterr()
assert "cloud" not in captured.err.lower()
class TestEmbeddingStoreModelPassthrough:
"""Tests that EmbeddingStore passes model to get_provider."""
def test_model_forwarded_to_get_provider(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None) as mock_gp:
EmbeddingStore(db, model="custom/model").close()
mock_gp.assert_called_once_with(None, model="custom/model")
def test_provider_and_model_forwarded(self, tmp_path):
db = tmp_path / "embeddings.db"
with patch("code_review_graph.embeddings.get_provider", return_value=None) as mock_gp:
EmbeddingStore(db, provider="local", model="custom/model").close()
mock_gp.assert_called_once_with("local", model="custom/model")
class TestMiniMaxEmbeddingProvider:
"""Unit tests for MiniMaxEmbeddingProvider."""
def test_name(self):
provider = MiniMaxEmbeddingProvider(api_key="test-key")
assert provider.name == "minimax:embo-01"
def test_dimension(self):
provider = MiniMaxEmbeddingProvider(api_key="test-key")
assert provider.dimension == 1536
def test_embed_calls_api_with_db_type(self):
provider = MiniMaxEmbeddingProvider(api_key="test-key")
mock_vectors = [[0.1] * 1536, [0.2] * 1536]
mock_response = json.dumps({
"vectors": mock_vectors,
"total_tokens": 10,
"base_resp": {"status_code": 0, "status_msg": "success"},
}).encode("utf-8")
mock_resp_obj = MagicMock()
mock_resp_obj.read.return_value = mock_response
mock_resp_obj.__enter__ = MagicMock(return_value=mock_resp_obj)
mock_resp_obj.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock_resp_obj) as mock_urlopen:
result = provider.embed(["hello", "world"])
assert len(result) == 2
assert len(result[0]) == 1536
call_args = mock_urlopen.call_args
req = call_args[0][0]
payload = json.loads(req.data.decode("utf-8"))
assert payload["type"] == "db"
assert payload["model"] == "embo-01"
def test_embed_query_calls_api_with_query_type(self):
provider = MiniMaxEmbeddingProvider(api_key="test-key")
mock_vectors = [[0.5] * 1536]
mock_response = json.dumps({
"vectors": mock_vectors,
"total_tokens": 5,
"base_resp": {"status_code": 0, "status_msg": "success"},
}).encode("utf-8")
mock_resp_obj = MagicMock()
mock_resp_obj.read.return_value = mock_response
mock_resp_obj.__enter__ = MagicMock(return_value=mock_resp_obj)
mock_resp_obj.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock_resp_obj) as mock_urlopen:
result = provider.embed_query("search term")
assert len(result) == 1536
call_args = mock_urlopen.call_args
req = call_args[0][0]
payload = json.loads(req.data.decode("utf-8"))
assert payload["type"] == "query"
def test_embed_api_error_raises(self):
provider = MiniMaxEmbeddingProvider(api_key="test-key")
mock_response = json.dumps({
"vectors": [],
"base_resp": {"status_code": 1001, "status_msg": "invalid api key"},
}).encode("utf-8")
mock_resp_obj = MagicMock()
mock_resp_obj.read.return_value = mock_response
mock_resp_obj.__enter__ = MagicMock(return_value=mock_resp_obj)
mock_resp_obj.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock_resp_obj):
with pytest.raises(RuntimeError, match="invalid api key"):
provider.embed_query("test")
def test_embed_sends_user_agent_header(self):
# urllib's default UA ("Python-urllib/X.Y") is rejected by some
# Cloudflare-fronted gateways with HTTP 403 / error 1010. CRG must
# send an explicit User-Agent so requests get through.
provider = MiniMaxEmbeddingProvider(api_key="test-key")
mock_response = json.dumps({
"vectors": [[0.1] * 1536],
"total_tokens": 1,
"base_resp": {"status_code": 0, "status_msg": "success"},
}).encode("utf-8")
mock_resp_obj = MagicMock()
mock_resp_obj.read.return_value = mock_response
mock_resp_obj.__enter__ = MagicMock(return_value=mock_resp_obj)
mock_resp_obj.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock_resp_obj) as mock_urlopen:
provider.embed_query("hello")
req = mock_urlopen.call_args[0][0]
ua = req.headers.get("User-agent", "")
assert ua.startswith("code-review-graph/")
assert "github.com/tirth8205/code-review-graph" in ua
class TestGetProviderMiniMax:
"""Tests for get_provider() with MiniMax."""
def test_get_provider_minimax_with_key(self):
with patch.dict("os.environ", {"MINIMAX_API_KEY": "test-key"}):
provider = get_provider("minimax")
assert isinstance(provider, MiniMaxEmbeddingProvider)
assert provider.name == "minimax:embo-01"
def test_get_provider_minimax_without_key_raises(self):
with patch.dict("os.environ", {}, clear=True):
with pytest.raises(ValueError, match="MINIMAX_API_KEY"):
get_provider("minimax")
def _make_voyage_response(vectors: list[list[float]]) -> MagicMock:
body = json.dumps({
"data": [{"embedding": v, "index": i} for i, v in enumerate(vectors)],
"model": "voyage-code-3",
"object": "list",
"usage": {"total_tokens": 5},
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = body
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
return mock
class TestVoyageEmbeddingProvider:
"""Unit tests for VoyageEmbeddingProvider."""
def test_name_includes_model_dimension_dtype_and_endpoint(self):
provider = VoyageEmbeddingProvider(
api_key="k",
base_url="https://api.voyageai.com/v1/",
model="voyage-code-3",
output_dimension=1024,
output_dtype="float",
)
assert (
provider.name
== "voyage:voyage-code-3:dim1024:float@https://api.voyageai.com/v1"
)
def test_default_dimension_before_call(self):
provider = VoyageEmbeddingProvider(api_key="k")
assert provider.dimension == 1024
def test_embed_calls_api_with_document_input_type(self):
provider = VoyageEmbeddingProvider(api_key="secret-key")
with patch(
"urllib.request.urlopen",
return_value=_make_voyage_response([[0.1] * 1024, [0.2] * 1024]),
) as mock_urlopen:
result = provider.embed(["hello", "world"])
assert len(result) == 2
assert len(result[0]) == 1024
req = mock_urlopen.call_args[0][0]
payload = json.loads(req.data.decode("utf-8"))
assert payload["model"] == "voyage-code-3"
assert payload["input"] == ["hello", "world"]
assert payload["input_type"] == "document"
assert payload["output_dimension"] == 1024
assert payload["output_dtype"] == "float"
assert req.headers["Authorization"] == "Bearer secret-key"
assert req.full_url == "https://api.voyageai.com/v1/embeddings"
def test_embed_query_calls_api_with_query_input_type(self):
provider = VoyageEmbeddingProvider(api_key="k")
with patch(
"urllib.request.urlopen",
return_value=_make_voyage_response([[0.5] * 1024]),
) as mock_urlopen:
result = provider.embed_query("search term")
assert len(result) == 1024
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert payload["input_type"] == "query"
def test_custom_output_dimension_and_dtype_forwarded(self):
provider = VoyageEmbeddingProvider(
api_key="k",
model="voyage-code-3",
output_dimension=512,
output_dtype="float",
)
with patch(
"urllib.request.urlopen",
return_value=_make_voyage_response([[0.1] * 512]),
) as mock_urlopen:
provider.embed_query("x")
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert payload["output_dimension"] == 512
assert provider.dimension == 512
def test_response_without_index_field_falls_back_to_server_order(self):
provider = VoyageEmbeddingProvider(api_key="k")
body = json.dumps({
"data": [
{"embedding": [1.0]},
{"embedding": [2.0]},
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = body
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
assert provider.embed(["a", "b"]) == [[1.0], [2.0]]
def test_response_length_mismatch_raises(self):
provider = VoyageEmbeddingProvider(api_key="k")
with patch(
"urllib.request.urlopen",
return_value=_make_voyage_response([[0.1]]),
):
with pytest.raises(RuntimeError, match="refusing to misalign"):
provider.embed(["a", "b"])
def test_min_interval_paces_batched_requests(self):
provider = VoyageEmbeddingProvider(
api_key="k",
batch_size=1,
min_interval_sec=2.0,
)
with (
patch(
"urllib.request.urlopen",
side_effect=[
_make_voyage_response([[0.1] * 1024]),
_make_voyage_response([[0.2] * 1024]),
],
),
patch(
"code_review_graph.embeddings.time.monotonic",
side_effect=[100.0, 101.0, 103.0],
),
patch("code_review_graph.embeddings.time.sleep") as sleep,
):
result = provider.embed(["hello", "world"])
assert len(result) == 2
sleep.assert_called_once_with(1.0)
class TestGetProviderVoyage:
"""Tests for get_provider() with Voyage."""
def test_get_provider_voyage_with_key(self):
env = {
"VOYAGE_API_KEY": "test-key",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}
with patch.dict(os.environ, env, clear=False):
provider = get_provider("voyage")
assert isinstance(provider, VoyageEmbeddingProvider)
assert provider.name == (
"voyage:voyage-code-3:dim1024:float@https://api.voyageai.com/v1"
)
def test_get_provider_voyage_without_key_raises(self):
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="VOYAGE_API_KEY"):
get_provider("voyage")
def test_get_provider_voyage_respects_env_configuration(self):
env = {
"VOYAGE_API_KEY": "test-key",
"CRG_VOYAGE_MODEL": "voyage-code-3",
"CRG_VOYAGE_BASE_URL": "https://voyage.example.test/v1",
"CRG_VOYAGE_OUTPUT_DIMENSION": "512",
"CRG_VOYAGE_OUTPUT_DTYPE": "float",
"CRG_VOYAGE_MIN_INTERVAL_SEC": "21",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}
with patch.dict(os.environ, env, clear=False):
provider = get_provider("voyage")
assert isinstance(provider, VoyageEmbeddingProvider)
assert provider._min_interval_sec == 21.0
assert provider.name == (
"voyage:voyage-code-3:dim512:float@https://voyage.example.test/v1"
)
def test_get_provider_voyage_ignores_local_embedding_model_env(self):
env = {
"VOYAGE_API_KEY": "test-key",
"CRG_EMBEDDING_MODEL": "local-only-model",
"CRG_ACCEPT_CLOUD_EMBEDDINGS": "1",
}
with patch.dict(os.environ, env, clear=False):
provider = get_provider("voyage")
assert isinstance(provider, VoyageEmbeddingProvider)
assert provider.name == (
"voyage:voyage-code-3:dim1024:float@https://api.voyageai.com/v1"
)
class TestEmbeddingStoreContextManager:
"""Regression tests for #260: EmbeddingStore must support the context
manager protocol so connections are cleaned up on exception."""
def test_supports_context_manager(self, tmp_path):
db = tmp_path / "embed_ctx.db"
with EmbeddingStore(db) as store:
assert store is not None
assert store.db_path == db
# After exiting, connection should be closed.
# (Attempting another query would fail, but we don't test that
# because close() doesn't invalidate the object — it just
# closes the underlying sqlite3 connection.)
def test_context_manager_closes_on_exception(self, tmp_path):
db = tmp_path / "embed_err.db"
try:
with EmbeddingStore(db) as store:
assert store.db_path == db
raise RuntimeError("simulated crash")
except RuntimeError:
pass
# The connection was closed by __exit__ even though an exception
# was raised. This is the whole point of #260 — without the
# context manager, the connection would leak.
def _make_openai_response(vectors: list[list[float]]) -> MagicMock:
body = json.dumps({
"data": [{"embedding": v, "index": i} for i, v in enumerate(vectors)],
"model": "text-embedding-3-small",
"object": "list",
"usage": {"prompt_tokens": 5, "total_tokens": 5},
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = body
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
return mock
@pytest.fixture
def openai_loopback_server():
"""Serve deterministic OpenAI-compatible embeddings on loopback."""
payloads: list[dict] = []
class Handler(BaseHTTPRequestHandler):
def do_POST(self):
size = int(self.headers["Content-Length"])
payload = json.loads(self.rfile.read(size))
payloads.append(payload)
dimension = payload.get("dimensions", 7)
response = {
"data": [
{"embedding": [0.1] * dimension, "index": index}
for index, _text in enumerate(payload["input"])
],
"model": payload["model"],
}
body = json.dumps(response).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, _format, *_args):
return
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
try:
host, port = server.server_address
yield f"http://{host}:{port}/v1", payloads
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
class TestIsLocalhostUrl:
"""Ensure localhost detection is robust against subdomain tricks."""
def test_plain_localhost(self):
assert _is_localhost_url("http://localhost:3000/v1")
def test_127_loopback(self):
assert _is_localhost_url("http://127.0.0.1:3000/v1")
def test_0000_loopback(self):
assert _is_localhost_url("http://0.0.0.0:8080/v1")
def test_ipv6_loopback(self):
assert _is_localhost_url("http://[::1]:3000/v1")
def test_real_cloud_host(self):
assert not _is_localhost_url("https://api.openai.com/v1")
def test_subdomain_spoof_not_localhost(self):
# Architect flagged: plain string match would mis-classify this.
assert not _is_localhost_url("https://my-openai.127.0.0.1.nip.io/v1")
def test_invalid_url(self):
assert not _is_localhost_url("not a url")
class TestOpenAIEmbeddingProvider:
def test_name_includes_model(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="text-embedding-3-small",
)
assert p.name == "openai:text-embedding-3-small@http://localhost:3000/v1"
def test_default_dimension_before_call(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
assert p.dimension == 1536 # fallback until first response
def test_dimension_captured_from_response(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 768]),
):
vec = p.embed_query("hello")
assert len(vec) == 768
assert p.dimension == 768
def test_embed_calls_api_with_correct_payload(self):
p = OpenAIEmbeddingProvider(
api_key="secret-key",
base_url="http://127.0.0.1:3000/v1",
model="text-embedding-3-small",
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 1536, [0.2] * 1536]),
) as mock_urlopen:
result = p.embed(["hello", "world"])
assert len(result) == 2
assert len(result[0]) == 1536
req = mock_urlopen.call_args[0][0]
payload = json.loads(req.data.decode("utf-8"))
assert payload["model"] == "text-embedding-3-small"
assert payload["input"] == ["hello", "world"]
assert "dimensions" not in payload # not pinned by default
assert req.headers["Authorization"] == "Bearer secret-key"
assert req.headers["Content-type"] == "application/json"
# Cloudflare-fronted gateways (e.g. Fireworks) reject the urllib
# default UA with HTTP 403 / error 1010. See _USER_AGENT in
# embeddings.py.
ua = req.headers.get("User-agent", "")
assert ua.startswith("code-review-graph/")
assert "github.com/tirth8205/code-review-graph" in ua
assert req.full_url == "http://127.0.0.1:3000/v1/embeddings"
def test_explicit_dimension_forwarded_in_payload(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1",
model="text-embedding-3-large", dimension=256,
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 256]),
) as mock_urlopen:
p.embed_query("x")
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert payload["dimensions"] == 256
def test_loopback_auto_learned_dimension_is_never_resent(
self, openai_loopback_server,
):
base_url, payloads = openai_loopback_server
provider = OpenAIEmbeddingProvider(
api_key="k",
base_url=base_url,
model="text-embedding-3-small",
)
assert len(provider.embed_query("first")) == 7
assert provider.dimension == 7
assert len(provider.embed_query("second")) == 7
assert len(payloads) == 2
assert all("dimensions" not in payload for payload in payloads)
def test_loopback_explicit_dimension_is_sent_for_custom_model_alias(
self, openai_loopback_server,
):
base_url, payloads = openai_loopback_server
provider = OpenAIEmbeddingProvider(
api_key="k",
base_url=base_url,
model="azure-production-deployment",
dimension=4,
)
assert len(provider.embed_query("custom alias")) == 4
assert payloads == [{
"model": "azure-production-deployment",
"input": ["custom alias"],
"dimensions": 4,
}]
def test_auto_learned_dimension_omitted_for_non_v3_models(self):
# Many OpenAI-compatible providers (SiliconFlow, Cohere, voyage-3,
# custom vLLM gateways) reject the `dimensions` body field with
# HTTP 400. The provider auto-learns dimension from the first
# response and would otherwise forward it on every subsequent call.
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1",
model="BAAI/bge-m3",
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 1024]),
) as mock_urlopen:
vec = p.embed_query("x")
assert len(vec) == 1024
assert p.dimension == 1024
# Second call would have auto-forwarded dimensions before the fix.
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 1024]),
) as mock_urlopen:
p.embed_query("y")
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert "dimensions" not in payload, (
f"non-v3 model {p._model!r} should not send `dimensions`; "
f"got payload keys: {list(payload)}"
)
def test_explicit_dimension_forwarded_for_non_v3_models(self):
# Explicit requests must not be inferred from the model name. OpenAI-
# compatible gateways can expose dimension-capable models under
# arbitrary aliases.
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1",
model="BAAI/bge-m3", dimension=1024,
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 1024]),
) as mock_urlopen:
p.embed_query("x")
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert payload["dimensions"] == 1024
def test_explicit_dimension_forwarded_for_v3_models(self):
# Regression guard: v3 models must still honor the pinned dimension.
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1",
model="text-embedding-3-large", dimension=512,
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 512]),
) as mock_urlopen:
p.embed_query("x")
payload = json.loads(mock_urlopen.call_args[0][0].data.decode("utf-8"))
assert payload["dimensions"] == 512
def test_base_url_trailing_slash_stripped(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1/", model="m",
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 10]),
) as mock_urlopen:
p.embed_query("x")
req = mock_urlopen.call_args[0][0]
assert req.full_url == "http://localhost:3000/v1/embeddings"
def test_embed_api_error_raises(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
err_body = json.dumps({
"error": {"message": "invalid api key", "type": "invalid_request_error"},
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = err_body
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="invalid api key"):
p.embed_query("x")
def test_embed_empty_data_raises(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
body = json.dumps({"data": []}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = body
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="empty data"):
p.embed_query("x")
def test_batching_splits_into_100_per_request(self):
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
texts = [f"text-{i}" for i in range(250)]
call_count = {"n": 0}
def _mk_response(*_args, **_kwargs):
call_count["n"] += 1
# match payload size
req = _args[0]
body = json.loads(req.data.decode("utf-8"))
n = len(body["input"])
return _make_openai_response([[0.1] * 5 for _ in range(n)])
with patch("urllib.request.urlopen", side_effect=_mk_response):
out = p.embed(texts)
assert len(out) == 250
assert call_count["n"] == 3 # 100 + 100 + 50
def test_custom_batch_size_respected(self):
"""new-api gateways (e.g. text-embedding-v4) cap batch at 10 —
user must be able to lower the batch size to avoid 400 errors."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
batch_size=10,
)
texts = [f"t-{i}" for i in range(25)]
call_count = {"n": 0}
def _mk_response(*_args, **_kwargs):
call_count["n"] += 1
req = _args[0]
body = json.loads(req.data.decode("utf-8"))
assert len(body["input"]) <= 10 # never exceed configured size
return _make_openai_response([[0.1] * 5 for _ in body["input"]])
with patch("urllib.request.urlopen", side_effect=_mk_response):
out = p.embed(texts)
assert len(out) == 25
assert call_count["n"] == 3 # 10 + 10 + 5
def test_empty_input_returns_empty(self):
"""embed([]) must short-circuit without hitting the API."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
with patch("urllib.request.urlopen") as mock_urlopen:
assert p.embed([]) == []
mock_urlopen.assert_not_called()
def test_endpoint_isolation_in_name(self):
"""Two providers with the same model but different base URLs MUST
produce different provider.name values, otherwise the embeddings
store silently reuses vectors from a different backend's vector space.
(Codex review HIGH finding.)"""
p1 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://api.openai.com/v1",
model="text-embedding-3-small",
)
p2 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://openrouter.ai/api/v1",
model="text-embedding-3-small",
)
p3 = OpenAIEmbeddingProvider(
api_key="k", base_url="http://127.0.0.1:3000/v1",
model="text-embedding-3-small",
)
assert p1.name != p2.name != p3.name
assert p1.name == "openai:text-embedding-3-small@https://api.openai.com/v1"
assert p2.name == "openai:text-embedding-3-small@https://openrouter.ai/api/v1"
assert p3.name == "openai:text-embedding-3-small@http://127.0.0.1:3000/v1"
def test_trailing_slash_does_not_change_identity(self):
"""A trailing slash on base_url must not cause a re-embed."""
p1 = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
p2 = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1/", model="m",
)
assert p1.name == p2.name
def test_path_routed_gateways_get_distinct_identity(self):
"""Path-routed gateways (same host, different URL path) front
different backends and must NOT share cached vectors.
(Codex round-2 HIGH finding.)"""
p1 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://gw.example.com/openai/v1", model="m",
)
p2 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://gw.example.com/vendor-b/v1", model="m",
)
assert p1.name != p2.name
assert p1.name == "openai:m@https://gw.example.com/openai/v1"
assert p2.name == "openai:m@https://gw.example.com/vendor-b/v1"
def test_default_port_is_stripped_from_identity(self):
"""`https://host/v1` and `https://host:443/v1` must map to the
same identity; stripping is necessary so the user can't force
a pointless re-embed by spelling the port differently.
(Codex round-2 MED finding.)"""
p1 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://api.openai.com/v1", model="m",
)
p2 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://api.openai.com:443/v1", model="m",
)
p3 = OpenAIEmbeddingProvider(
api_key="k", base_url="http://example.com:80/v1", model="m",
)
p4 = OpenAIEmbeddingProvider(
api_key="k", base_url="http://example.com/v1", model="m",
)
assert p1.name == p2.name
assert p3.name == p4.name
# Non-default port still affects identity (normal case).
p5 = OpenAIEmbeddingProvider(
api_key="k", base_url="https://api.openai.com:8443/v1", model="m",
)
assert p5.name != p1.name
def test_userinfo_is_stripped_from_identity(self):
"""Credentials embedded in the URL must NOT appear in provider.name
(which gets persisted into the embeddings table). This is an
at-rest credential-leak defense. (Codex round-2 MED finding.)"""
p_plain = OpenAIEmbeddingProvider(
api_key="k", base_url="https://api.example.com/v1", model="m",
)
p_auth = OpenAIEmbeddingProvider(
api_key="k", base_url="https://user:[email protected]/v1", model="m",
)
# 1. Same identity — userinfo stripped.
assert p_plain.name == p_auth.name
# 2. The secret never appears in the identity string.
assert "secret" not in p_auth.name
assert "user" not in p_auth.name
def test_ipv6_literal_in_identity(self):
"""IPv6 hostnames must round-trip cleanly, with brackets restored
when a non-default port is attached."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://[::1]:3000/v1", model="m",
)
assert p.name == "openai:m@http://[::1]:3000/v1"
def test_response_with_missing_index_raises(self):
"""Length-only checks let duplicate/missing indices through. We
require a strict 0..N-1 permutation. (Codex round-2 MED finding.)"""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
bad = json.dumps({
"data": [
{"embedding": [1.0], "index": 0},
{"embedding": [2.0], "index": 0}, # duplicate 0, missing 1
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = bad
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="malformed indices"):
p.embed(["a", "b"])
def test_response_with_out_of_range_index_raises(self):
"""Index >= N is invalid even if count matches."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
bad = json.dumps({
"data": [
{"embedding": [1.0], "index": 0},
{"embedding": [2.0], "index": 5}, # out-of-range
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = bad
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="malformed indices"):
p.embed(["a", "b"])
def test_response_without_index_field_falls_back_to_server_order(self):
"""Some OpenAI-compatible gateways omit `index` entirely. The
length check is the only safety net available — we must still
succeed on length match and fail on mismatch."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
no_idx = json.dumps({
"data": [
{"embedding": [1.0]},
{"embedding": [2.0]},
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = no_idx
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
result = p.embed(["a", "b"])
# Trust server order when index is absent.
assert result == [[1.0], [2.0]]
def test_scheme_change_produces_distinct_identity(self):
"""http and https to the same host/path front different endpoints
in practice (dev vs prod gateway, pre/post TLS migration). They
must NOT share cached vectors. (Codex round-3 HIGH finding.)"""
p_http = OpenAIEmbeddingProvider(
api_key="k", base_url="http://gw.example.com/v1", model="m",
)
p_https = OpenAIEmbeddingProvider(
api_key="k", base_url="https://gw.example.com/v1", model="m",
)
assert p_http.name != p_https.name
# http default port 80 and https default port 443 are both stripped
# from the host, but scheme is preserved in the identity.
assert p_http.name == "openai:m@http://gw.example.com/v1"
assert p_https.name == "openai:m@https://gw.example.com/v1"
def test_mixed_indexed_unindexed_response_raises(self):
"""Some items with ``index``, others without: must refuse rather
than silently zip in server order (which would misplace the
indexed items). (Codex round-3 HIGH finding.)"""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
mixed = json.dumps({
"data": [
{"embedding": [1.0], "index": 1}, # claims to be for input[1]
{"embedding": [2.0]}, # no index
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = mixed
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="mixed indexed/unindexed"):
p.embed(["a", "b"])
def test_string_index_treated_as_mixed(self):
"""Some OpenAI-compatible gateways serialize index as a string.
Our permutation check requires ints; string index must fall to
the mixed-case refusal, not silently slip through."""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
bad = json.dumps({
"data": [
{"embedding": [1.0], "index": "0"}, # string, not int
{"embedding": [2.0], "index": "1"},
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = bad
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
with pytest.raises(RuntimeError, match="mixed indexed/unindexed"):
p.embed(["a", "b"])
def test_retry_on_remote_disconnected(self, monkeypatch):
"""http.client.RemoteDisconnected is a common transient failure
when reverse proxies drop idle connections. Must retry.
(Codex round-2 LOW finding.)"""
import http.client
monkeypatch.setattr(time, "sleep", lambda s: None)
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
call_count = {"n": 0}
def _mock_urlopen(*args, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
raise http.client.RemoteDisconnected("edge proxy dropped connection")
return _make_openai_response([[0.1] * 5])
with patch("urllib.request.urlopen", side_effect=_mock_urlopen):
p.embed_query("x")
assert call_count["n"] == 2
def test_response_length_mismatch_raises(self):
"""Gateway returns fewer embeddings than inputs: refuse to proceed
rather than silently zip misaligned vectors onto the wrong nodes.
(Codex review MED finding.)"""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
with patch(
"urllib.request.urlopen",
return_value=_make_openai_response([[0.1] * 5]), # 1 vec
):
with pytest.raises(RuntimeError, match="refusing to misalign"):
p.embed(["a", "b", "c"]) # 3 inputs
def test_reordered_response_is_sorted_by_index(self):
"""Gateway returns data out of order: restore input order via
the `index` field, so vec[i] always corresponds to input[i].
(Codex review MED finding.)"""
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
# Return data in order 2, 0, 1 (i.e. reversed-ish).
reordered = json.dumps({
"data": [
{"embedding": [3.0], "index": 2},
{"embedding": [1.0], "index": 0},
{"embedding": [2.0], "index": 1},
],
}).encode("utf-8")
mock = MagicMock()
mock.read.return_value = reordered
mock.__enter__ = MagicMock(return_value=mock)
mock.__exit__ = MagicMock(return_value=False)
with patch("urllib.request.urlopen", return_value=mock):
result = p.embed(["a", "b", "c"])
# Must be [[1.0], [2.0], [3.0]] after sorting by index.
assert result == [[1.0], [2.0], [3.0]]
def test_retry_on_http_429(self, monkeypatch):
"""HTTP 429 must trigger retry with backoff (not bail immediately).
(Codex review MED finding — prior substring match missed the fact
that error bodies may not contain '429'.)"""
import urllib.error
monkeypatch.setattr(time, "sleep", lambda s: None) # instant retries
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
call_count = {"n": 0}
good_response = _make_openai_response([[0.1] * 5])
import io
def _mock_urlopen(*args, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
raise urllib.error.HTTPError(
url="http://localhost:3000/v1/embeddings",
code=429, msg="Too Many Requests", hdrs=None,
fp=io.BytesIO(b'{"error": "rate limited"}'),
)
return good_response
with patch("urllib.request.urlopen", side_effect=_mock_urlopen):
out = p.embed_query("x")
assert len(out) == 5
assert call_count["n"] == 2 # 1 fail + 1 success
def test_retry_on_socket_timeout(self, monkeypatch):
"""socket.timeout (read timeout) must be classified retryable —
previously these surfaced as str(exc) without '429/500/503' so
retry never fired. (Codex review MED finding.)"""
import socket
monkeypatch.setattr(time, "sleep", lambda s: None)
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
call_count = {"n": 0}
good_response = _make_openai_response([[0.1] * 5])
def _mock_urlopen(*args, **kwargs):
call_count["n"] += 1
if call_count["n"] <= 2:
raise socket.timeout("read timed out")
return good_response
with patch("urllib.request.urlopen", side_effect=_mock_urlopen):
out = p.embed_query("x")
assert len(out) == 5
assert call_count["n"] == 3 # 2 fails + 1 success
def test_retry_on_url_error(self, monkeypatch):
"""URLError (connection refused, DNS failure) must retry."""
import urllib.error
monkeypatch.setattr(time, "sleep", lambda s: None)
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
call_count = {"n": 0}
def _mock_urlopen(*args, **kwargs):
call_count["n"] += 1
if call_count["n"] == 1:
raise urllib.error.URLError("connection refused")
return _make_openai_response([[0.1] * 5])
with patch("urllib.request.urlopen", side_effect=_mock_urlopen):
p.embed_query("x")
assert call_count["n"] == 2
def test_no_retry_on_http_400(self, monkeypatch):
"""HTTP 400 = caller bug (bad payload, unsupported model). Must fail
fast rather than waste time on 3 retries."""
import io
import urllib.error
monkeypatch.setattr(time, "sleep", lambda s: None)
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
call_count = {"n": 0}
def _mock_urlopen(*args, **kwargs):
call_count["n"] += 1
raise urllib.error.HTTPError(
url="http://localhost:3000/v1/embeddings",
code=400, msg="Bad Request", hdrs=None,
fp=io.BytesIO(b'{"error": {"message": "invalid model"}}'),
)
with patch("urllib.request.urlopen", side_effect=_mock_urlopen):
with pytest.raises(RuntimeError, match="invalid model"):
p.embed_query("x")
assert call_count["n"] == 1 # no retry on 4xx non-429
def test_http_error_body_is_surfaced(self):
"""If the gateway returns 400 with a JSON error body, the RuntimeError
must include the real reason, not just 'HTTP Error 400: Bad Request'."""
import urllib.error
p = OpenAIEmbeddingProvider(
api_key="k", base_url="http://localhost:3000/v1", model="m",
)
body = json.dumps({
"error": {"message": "batch size is invalid, should not exceed 10."},
}).encode("utf-8")
# HTTPError's .read() returns bytes from its fp
import io
err = urllib.error.HTTPError(
url="http://localhost:3000/v1/embeddings",
code=400, msg="Bad Request", hdrs=None, fp=io.BytesIO(body),
)
with patch("urllib.request.urlopen", side_effect=err):
with pytest.raises(RuntimeError, match="batch size is invalid"):
p.embed_query("x")
class TestGetProviderOpenAI:
_MIN_ENV = {
"CRG_OPENAI_API_KEY": "sk-test",
"CRG_OPENAI_BASE_URL": "http://127.0.0.1:3000/v1",
"CRG_OPENAI_MODEL": "text-embedding-3-small",
}
def test_with_all_env_vars(self):
with patch.dict("os.environ", self._MIN_ENV, clear=True):
p = get_provider("openai")
assert isinstance(p, OpenAIEmbeddingProvider)
assert p.name == "openai:text-embedding-3-small@http://127.0.0.1:3000/v1"
def test_missing_api_key_raises(self):
env = {k: v for k, v in self._MIN_ENV.items() if k != "CRG_OPENAI_API_KEY"}
with patch.dict("os.environ", env, clear=True):
with pytest.raises(ValueError, match="CRG_OPENAI_API_KEY"):
get_provider("openai")
def test_missing_base_url_raises(self):
env = {k: v for k, v in self._MIN_ENV.items() if k != "CRG_OPENAI_BASE_URL"}
with patch.dict("os.environ", env, clear=True):
with pytest.raises(ValueError, match="CRG_OPENAI_BASE_URL"):
get_provider("openai")
def test_missing_model_raises(self):
env = {k: v for k, v in self._MIN_ENV.items() if k != "CRG_OPENAI_MODEL"}
with patch.dict("os.environ", env, clear=True):
with pytest.raises(ValueError, match="CRG_OPENAI_MODEL"):
get_provider("openai")
def test_model_arg_overrides_env(self):
with patch.dict("os.environ", self._MIN_ENV, clear=True):
p = get_provider("openai", model="text-embedding-3-large")
assert p.name == "openai:text-embedding-3-large@http://127.0.0.1:3000/v1"
def test_dimension_env_forwarded(self):
env = {**self._MIN_ENV, "CRG_OPENAI_DIMENSION": "256"}
with patch.dict("os.environ", env, clear=True):
p = get_provider("openai")
assert p._dimension == 256
def test_localhost_suppresses_egress_warning(self, capsys):
with patch.dict("os.environ", self._MIN_ENV, clear=True):
get_provider("openai")
captured = capsys.readouterr()
# localhost must never trigger the cloud-egress warning
assert captured.err == ""
assert captured.out == ""
def test_cloud_base_url_triggers_egress_warning(self, capsys):
env = {**self._MIN_ENV, "CRG_OPENAI_BASE_URL": "https://api.openai.com/v1"}
with patch.dict("os.environ", env, clear=True):
# drop accept flag to ensure warning fires
os.environ.pop("CRG_ACCEPT_CLOUD_EMBEDDINGS", None)
get_provider("openai")
captured = capsys.readouterr()
assert "openai" in captured.err.lower()
assert "cloud" in captured.err.lower()
assert captured.out == "" # MCP stdio safety
def test_subdomain_spoof_triggers_warning(self, capsys):
"""my-openai.127.0.0.1.nip.io must NOT be treated as localhost."""
env = {
**self._MIN_ENV,
"CRG_OPENAI_BASE_URL": "https://my-openai.127.0.0.1.nip.io/v1",
}
with patch.dict("os.environ", env, clear=True):
get_provider("openai")
captured = capsys.readouterr()
assert "cloud" in captured.err.lower()
Reference in New Issue View Git Blame Copy Permalink