switch to a local mock that only affects the one test

Author: miguelgrinberg

test_elasticsearch/test_dsl/conftest.py

@@ -19,11 +19,9 @@
 import asyncio
 import os
 import re
-import sys
 import time
 from datetime import datetime
-from hashlib import md5
-from typing import Any, AsyncGenerator, Dict, Generator, List, Tuple, cast
+from typing import Any, AsyncGenerator, Dict, Generator, Tuple, cast
 from unittest import SkipTest
 from unittest.mock import AsyncMock, Mock
 
@@ -52,54 +50,6 @@
 )
 
 
-def pytest_configure(config: "pytest.Config"):
-    # mock sentence-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

@@ -15,26 +15,51 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from typing import Any, Tuple
+import sys
+from hashlib import md5
+from typing import Any, List, Tuple
 from unittest import SkipTest
+from unittest.mock import Mock, patch
 
 import pytest
 
 from elasticsearch import AsyncElasticsearch
 
-from ..async_examples import vectors
-
 
 @pytest.mark.asyncio
 async def test_vector_search(
-    async_write_client: AsyncElasticsearch, es_version: Tuple[int, ...], mocker: Any
+    async_write_client: AsyncElasticsearch, es_version: Tuple[int, ...]
 ) -> None:
     # this test only runs on Elasticsearch >= 8.11 because the example uses
     # a dense vector without specifying an explicit size
     if es_version < (8, 11):
         raise SkipTest("This test requires Elasticsearch 8.11 or newer")
 
-    await vectors.create()
-    await vectors.WorkplaceDoc._index.refresh()
-    results = await (await vectors.search("Welcome to our team!")).execute()
-    assert results[0].name == "Intellectual Property Policy"
+    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]
+
+    def mock_nltk_tokenize(content: str):
+        return content.split("\n")
+
+    # mock sentence_transformare and nltk, because they are quite big and
+    # irrelevant for testing the example logic
+    with patch.dict(
+        sys.modules,
+        {
+            "sentence_transformers": Mock(SentenceTransformer=MockSentenceTransformer),
+            "nltk": Mock(sent_tokenize=mock_nltk_tokenize),
+        },
+    ):
+        # import the example after the dependencies are mocked
+        from ..async_examples import vectors
+
+        await vectors.create()
+        await vectors.WorkplaceDoc._index.refresh()
+        results = await (await vectors.search("Welcome to our team!")).execute()
+        assert results[0].name == "Intellectual Property Policy"

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

@@ -15,26 +15,51 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
-from typing import Any, Tuple
+import sys
+from hashlib import md5
+from typing import Any, List, Tuple
 from unittest import SkipTest
+from unittest.mock import Mock, patch
 
 import pytest
 
 from elasticsearch import Elasticsearch
 
-from ..examples import vectors
-
 
 @pytest.mark.sync
 def test_vector_search(
-    write_client: Elasticsearch, es_version: Tuple[int, ...], mocker: Any
+    write_client: Elasticsearch, es_version: Tuple[int, ...]
 ) -> None:
     # this test only runs on Elasticsearch >= 8.11 because the example uses
     # a dense vector without specifying an explicit size
     if es_version < (8, 11):
         raise SkipTest("This test requires Elasticsearch 8.11 or newer")
 
-    vectors.create()
-    vectors.WorkplaceDoc._index.refresh()
-    results = (vectors.search("Welcome to our team!")).execute()
-    assert results[0].name == "Intellectual Property Policy"
+    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]
+
+    def mock_nltk_tokenize(content: str):
+        return content.split("\n")
+
+    # mock sentence_transformare and nltk, because they are quite big and
+    # irrelevant for testing the example logic
+    with patch.dict(
+        sys.modules,
+        {
+            "sentence_transformers": Mock(SentenceTransformer=MockSentenceTransformer),
+            "nltk": Mock(sent_tokenize=mock_nltk_tokenize),
+        },
+    ):
+        # import the example after the dependencies are mocked
+        from ..examples import vectors
+
+        vectors.create()
+        vectors.WorkplaceDoc._index.refresh()
+        results = (vectors.search("Welcome to our team!")).execute()
+        assert results[0].name == "Intellectual Property Policy"