Nucliadb test unification: search (part 1)

nucliadb/tests/conftest.py

@@ -36,6 +36,7 @@
         "tests.ndbfixtures.standalone",
         "tests.ndbfixtures.reader",
         "tests.ndbfixtures.writer",
+        "tests.ndbfixtures.search",
         "tests.ndbfixtures.train",
         "tests.ndbfixtures.ingest",
         # subcomponents

nucliadb/tests/fixtures.py

@@ -352,6 +352,8 @@ async def knowledge_graph(nucliadb_writer: AsyncClient, nucliadb_grpc: WriterStu
     return (nodes, edges)
 
 
+# TODO: remove after migrating tests/nucliadb/ to ndbfixtures. fixture already
+# moved to ndbfixtures.common
 @pytest.fixture(scope="function")
 async def stream_audit(natsd: str, mocker):
     from nucliadb_utils.audit.stream import StreamAuditStorage

nucliadb/tests/ingest/fixtures.py

@@ -23,7 +23,7 @@
 from dataclasses import dataclass
 from datetime import datetime
 from os.path import dirname, getsize
-from typing import AsyncIterator, Iterable, Iterator, Optional
+from typing import AsyncIterator, Iterable, Optional
 from unittest.mock import AsyncMock, patch
 
 import nats
@@ -49,14 +49,11 @@
 from nucliadb_protos.knowledgebox_pb2 import SemanticModelMetadata
 from nucliadb_protos.writer_pb2 import BrokerMessage
 from nucliadb_utils import const
-from nucliadb_utils.audit.audit import AuditStorage
-from nucliadb_utils.audit.basic import BasicAuditStorage
-from nucliadb_utils.audit.stream import StreamAuditStorage
 from nucliadb_utils.cache.nats import NatsPubsub
 from nucliadb_utils.cache.pubsub import PubSubDriver
 from nucliadb_utils.indexing import IndexingUtility
 from nucliadb_utils.nats import NatsConnectionManager
-from nucliadb_utils.settings import audit_settings, indexing_settings, transaction_settings
+from nucliadb_utils.settings import indexing_settings, transaction_settings
 from nucliadb_utils.storages.settings import settings as storage_settings
 from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.transaction import TransactionUtility
@@ -225,33 +222,6 @@ async def knowledgebox_with_vectorsets(
     await KnowledgeBox.delete(maindb_driver, kbid)
 
 
-@pytest.fixture(scope="function")
-def audit(basic_audit: BasicAuditStorage) -> Iterator[AuditStorage]:
-    # XXX: why aren't we settings the utility?
-    yield basic_audit
-
-
-@pytest.fixture(scope="function")
-async def basic_audit() -> AsyncIterator[BasicAuditStorage]:
-    audit = BasicAuditStorage()
-    await audit.initialize()
-    yield audit
-    await audit.finalize()
-
-
-@pytest.fixture(scope="function")
-async def stream_audit(nats_server: str) -> AsyncIterator[StreamAuditStorage]:
-    audit = StreamAuditStorage(
-        [nats_server],
-        audit_settings.audit_jetstream_target,  # type: ignore
-        audit_settings.audit_partitions,
-        audit_settings.audit_hash_seed,
-    )
-    await audit.initialize()
-    yield audit
-    await audit.finalize()
-
-
 @pytest.fixture(scope="function")
 async def indexing_utility(
     dummy_indexing_utility: IndexingUtility,

nucliadb/tests/ndbfixtures/common.py

@@ -18,31 +18,76 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from os.path import dirname
-from typing import AsyncIterator
+from typing import AsyncIterator, Iterator
 
 import pytest
+from pytest_mock import MockerFixture
 
 from nucliadb.common.cluster.manager import KBShardManager
 from nucliadb.common.maindb.driver import Driver
+from nucliadb_utils.audit.audit import AuditStorage
+from nucliadb_utils.audit.basic import BasicAuditStorage
+from nucliadb_utils.audit.stream import StreamAuditStorage
+from nucliadb_utils.settings import audit_settings
 from nucliadb_utils.storages.settings import settings as storage_settings
 from nucliadb_utils.storages.storage import Storage
-from nucliadb_utils.utilities import (
-    Utility,
-    clean_utility,
-    set_utility,
-)
+from nucliadb_utils.utilities import Utility, clean_utility, set_utility
+from tests.ndbfixtures.utils import global_utility
+
+# Audit
 
 
 @pytest.fixture(scope="function")
-async def shard_manager(storage: Storage, maindb_driver: Driver) -> AsyncIterator[KBShardManager]:
-    sm = KBShardManager()
-    set_utility(Utility.SHARD_MANAGER, sm)
+def audit(basic_audit: BasicAuditStorage) -> Iterator[AuditStorage]:
+    yield basic_audit
 
-    yield sm
 
-    clean_utility(Utility.SHARD_MANAGER)
+@pytest.fixture(scope="function")
+async def basic_audit() -> AsyncIterator[BasicAuditStorage]:
+    audit = BasicAuditStorage()
+    await audit.initialize()
+    with global_utility(Utility.AUDIT, audit):
+        yield audit
+    await audit.finalize()
+
+
+@pytest.fixture(scope="function")
+async def stream_audit(nats_server: str, mocker: MockerFixture) -> AsyncIterator[StreamAuditStorage]:
+    audit = StreamAuditStorage(
+        [nats_server],
+        audit_settings.audit_jetstream_target,  # type: ignore
+        audit_settings.audit_partitions,
+        audit_settings.audit_hash_seed,
+    )
+    await audit.initialize()
+
+    mocker.spy(audit, "send")
+    mocker.spy(audit.js, "publish")
+    mocker.spy(audit, "search")
+    mocker.spy(audit, "chat")
+
+    with global_utility(Utility.AUDIT, audit):
+        yield audit
+
+    await audit.finalize()
+
+
+# Local files
 
 
 @pytest.fixture(scope="function")
 async def local_files():
     storage_settings.local_testing_files = f"{dirname(__file__)}"
+
+
+# Shard manager
+
+
+@pytest.fixture(scope="function")
+async def shard_manager(storage: Storage, maindb_driver: Driver) -> AsyncIterator[KBShardManager]:
+    sm = KBShardManager()
+    set_utility(Utility.SHARD_MANAGER, sm)
+
+    yield sm
+
+    clean_utility(Utility.SHARD_MANAGER)