참고소스 수정본

참고/neo4j-graphrag-python-main/tests/unit/embeddings/test_sentence_transformers.py

@@ -0,0 +1,77 @@
+from unittest.mock import MagicMock, Mock, patch
+
+import numpy as np
+import pytest
+import torch
+from neo4j_graphrag.embeddings.base import Embedder
+from neo4j_graphrag.embeddings.sentence_transformers import (
+    SentenceTransformerEmbeddings,
+)
+from neo4j_graphrag.exceptions import EmbeddingsGenerationError
+
+
+def get_mock_sentence_transformers() -> MagicMock:
+    mock = MagicMock()
+    # I know, I know... ¯\_(ツ)_/¯
+    # This is to cover the if type checks in the embed_query method
+    mock.Tensor = torch.Tensor
+    mock.ndarray = np.ndarray
+    return mock
+
+
+@patch("builtins.__import__")
+def test_initialization(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    instance = SentenceTransformerEmbeddings()
+    MockSentenceTransformer.SentenceTransformer.assert_called_with("all-MiniLM-L6-v2")
+    assert isinstance(instance, Embedder)
+
+
+@patch("builtins.__import__")
+def test_initialization_with_custom_model(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    custom_model = "distilbert-base-nli-stsb-mean-tokens"
+    SentenceTransformerEmbeddings(model=custom_model)
+    MockSentenceTransformer.SentenceTransformer.assert_called_with(custom_model)
+
+
+@patch("builtins.__import__")
+def test_embed_query(mock_import: Mock) -> None:
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    mock_model = MockSentenceTransformer.SentenceTransformer.return_value
+    mock_model.encode.return_value = np.array([[0.1, 0.2, 0.3]])
+
+    instance = SentenceTransformerEmbeddings()
+    result = instance.embed_query("test query")
+
+    mock_model.encode.assert_called_with(["test query"])
+    assert isinstance(result, list)
+    assert result == [0.1, 0.2, 0.3]
+
+
+@patch("builtins.__import__", side_effect=ImportError)
+def test_import_error(mock_import: Mock) -> None:
+    with pytest.raises(ImportError):
+        SentenceTransformerEmbeddings()
+
+
+@patch("builtins.__import__")
+def test_embed_query_non_retryable_error_handling(mock_import: Mock) -> None:
+    """Test that non-retryable errors fail immediately without retries."""
+    MockSentenceTransformer = get_mock_sentence_transformers()
+    mock_import.return_value = MockSentenceTransformer
+    mock_model = MockSentenceTransformer.SentenceTransformer.return_value
+    mock_model.encode.side_effect = Exception("Model error")
+
+    instance = SentenceTransformerEmbeddings()
+    with pytest.raises(
+        EmbeddingsGenerationError,
+        match="Failed to generate embedding with SentenceTransformer",
+    ):
+        instance.embed_query("test query")
+
+    # Verify the model was called only once (no retries for non-rate-limit errors)
+    assert mock_model.encode.call_count == 1