Mock sentence-transformers and nltk in tests

Author: miguelgrinberg

pyproject.toml

@@ -77,8 +77,6 @@ dev = [
     "pandas",
     "mapbox-vector-tile",
     "jinja2",
-    "nltk",
-    "sentence_transformers",
     "tqdm",
     "mypy",
     "pyright",

test_elasticsearch/test_dsl/conftest.py

@@ -19,12 +19,15 @@
 import asyncio
 import os
 import re
+import sys
 import time
 from datetime import datetime
-from typing import Any, AsyncGenerator, Dict, Generator, Tuple, cast
+from hashlib import md5
+from typing import Any, AsyncGenerator, Dict, Generator, List, Tuple, cast
 from unittest import SkipTest
 from unittest.mock import AsyncMock, Mock
 
+import pytest
 import pytest_asyncio
 from elastic_transport import ObjectApiResponse
 from pytest import fixture, skip
@@ -49,6 +52,54 @@
 )
 
 
+def pytest_configure(config: "pytest.Config"):
+    # mock setence-transformers
+    class MockSentenceTransformer:
+        def __init__(self, model: Any):
+            pass
+
+        def encode(self, text: str) -> List[float]:
+            vector = [int(ch) for ch in md5(text.encode()).digest()]
+            total = sum(vector)
+            return [float(v) / total for v in vector]
+
+    mock_sentence_transformers_mod = type(sys)("sentence_transformers")
+    setattr(
+        mock_sentence_transformers_mod,
+        "__original_mod",
+        sys.modules.get("sentence_transformers"),
+    )
+    setattr(
+        mock_sentence_transformers_mod, "SentenceTransformer", MockSentenceTransformer
+    )
+    sys.modules[mock_sentence_transformers_mod.__name__] = (
+        mock_sentence_transformers_mod
+    )
+
+    # mock nltk
+    def mock_tokenize(content: str):
+        return content.split("\n")
+
+    mock_nltk_mod = type(sys)("nlkt")
+    setattr(mock_nltk_mod, "__original_mod", sys.modules.get("nltk"))
+    setattr(mock_nltk_mod, "download", Mock())
+    setattr(mock_nltk_mod, "sent_tokenize", mock_tokenize)
+    sys.modules["nltk"] = mock_nltk_mod
+
+
+def pytest_unconfigure(config: "pytest.Config"):
+    original_sentence_transformers = sys.modules["sentence_transformers"].__original_mod
+    if original_sentence_transformers:
+        sys.modules["sentence_transformers"] = original_sentence_transformers
+    else:
+        del sys.modules["sentence_transformers"]
+    original_nltk = sys.modules["nltk"].__original_mod
+    if original_nltk:
+        sys.modules["nltk"] = original_nltk
+    else:
+        del sys.modules["nltk"]
+
+
 def get_test_client(
     elasticsearch_url, wait: bool = True, **kwargs: Any
 ) -> Elasticsearch:

test_elasticsearch/test_dsl/test_integration/test_examples/_async/test_vectors.py

@@ -35,18 +35,7 @@ async def test_vector_search(
     if es_version < (8, 11):
         raise SkipTest("This test requires Elasticsearch 8.11 or newer")
 
-    class MockModel:
-        def __init__(self, model: Any):
-            pass
-
-        def encode(self, text: str) -> List[float]:
-            vector = [int(ch) for ch in md5(text.encode()).digest()]
-            total = sum(vector)
-            return [float(v) / total for v in vector]
-
-    mocker.patch.object(vectors, "SentenceTransformer", new=MockModel)
-
     await vectors.create()
     await vectors.WorkplaceDoc._index.refresh()
     results = await (await vectors.search("Welcome to our team!")).execute()
-    assert results[0].name == "New Employee Onboarding Guide"
+    assert results[0].name == "Intellectual Property Policy"

test_elasticsearch/test_dsl/test_integration/test_examples/_sync/test_vectors.py

@@ -35,18 +35,7 @@ def test_vector_search(
     if es_version < (8, 11):
         raise SkipTest("This test requires Elasticsearch 8.11 or newer")
 
-    class MockModel:
-        def __init__(self, model: Any):
-            pass
-
-        def encode(self, text: str) -> List[float]:
-            vector = [int(ch) for ch in md5(text.encode()).digest()]
-            total = sum(vector)
-            return [float(v) / total for v in vector]
-
-    mocker.patch.object(vectors, "SentenceTransformer", new=MockModel)
-
     vectors.create()
     vectors.WorkplaceDoc._index.refresh()
     results = (vectors.search("Welcome to our team!")).execute()
-    assert results[0].name == "New Employee Onboarding Guide"
+    assert results[0].name == "Intellectual Property Policy"