nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/ingest/tests/integration/consumer/__init__.py

@@ -1,18 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.

nucliadb/ingest/tests/integration/consumer/test_auditing.py

@@ -1,78 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-#
-
-import asyncio
-
-import pytest
-
-from nucliadb.ingest.consumer import auditing
-from nucliadb_protos import audit_pb2, writer_pb2
-from nucliadb_utils import const
-from nucliadb_utils.audit.stream import StreamAuditStorage
-
-pytestmark = pytest.mark.asyncio
-
-
-async def test_audit_counters(
-    stream_audit: StreamAuditStorage,
-    pubsub,
-    nats_manager,
-    fake_node,
-    knowledgebox_ingest,
-):
-    from nucliadb_utils.settings import audit_settings
-
-    partition = stream_audit.get_partition(knowledgebox_ingest)
-    subject = audit_settings.audit_jetstream_target.format(  # type: ignore
-        partition=partition, type="*"
-    )
-    await nats_manager.js.add_stream(
-        name=audit_settings.audit_stream, subjects=[subject]
-    )
-    psub = await nats_manager.js.pull_subscribe(subject, "psub")
-
-    iah = auditing.IndexAuditHandler(
-        audit=stream_audit,
-        pubsub=pubsub,
-        check_delay=0.05,
-    )
-    await iah.initialize()
-
-    await pubsub.publish(
-        const.PubSubChannels.RESOURCE_NOTIFY.format(kbid=knowledgebox_ingest),
-        writer_pb2.Notification(
-            kbid=knowledgebox_ingest,
-            action=writer_pb2.Notification.Action.INDEXED,
-        ).SerializeToString(),
-    )
-
-    await asyncio.sleep(0.1)
-
-    await iah.finalize()
-
-    # should have produced audit message, get the message
-    msg = await psub.fetch(1)
-    auditreq = audit_pb2.AuditRequest()
-    auditreq.ParseFromString(msg[0].data)
-
-    assert auditreq.kbid == knowledgebox_ingest
-    assert auditreq.type == audit_pb2.AuditRequest.AuditType.INDEXED
-    assert auditreq.kb_counter.fields == 2
-    assert auditreq.kb_counter.paragraphs == 2

nucliadb/ingest/tests/integration/consumer/test_materializer.py

@@ -1,126 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-#
-
-import asyncio
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-
-from nucliadb.common import datamanagers
-from nucliadb.ingest.consumer import materializer
-from nucliadb.ingest.tests.fixtures import create_resource
-from nucliadb_protos import writer_pb2
-from nucliadb_utils import const
-from nucliadb_utils.audit.stream import StreamAuditStorage
-from nucliadb_utils.nuclia_usage.protos.kb_usage_pb2 import KbUsage, Service
-from nucliadb_utils.utilities import Utility, clean_utility, set_utility
-
-pytestmark = pytest.mark.asyncio
-
-
-@pytest.fixture()
-def nats():
-    mock = AsyncMock()
-    mock.jetstream = MagicMock(return_value=AsyncMock())
-    yield mock
-
-
-@pytest.fixture()
-async def audit_storage(nats):
-    with patch("nucliadb_utils.audit.stream.nats.connect", return_value=nats):
-        aud = StreamAuditStorage(
-            nats_servers=["nats://localhost:4222"],
-            nats_target="test",
-            partitions=1,
-            seed=1,
-            nats_creds="nats_creds",
-        )
-        await aud.initialize()
-        set_utility(Utility.AUDIT, aud)
-        yield aud
-        clean_utility(Utility.AUDIT)
-        await aud.finalize()
-
-
-async def test_materialize_kb_data(
-    maindb_driver,
-    pubsub,
-    storage,
-    fake_node,
-    knowledgebox_ingest,
-    audit_storage,
-):
-    count = 10
-    for _ in range(count):
-        await create_resource(
-            storage=storage,
-            driver=maindb_driver,
-            knowledgebox_ingest=knowledgebox_ingest,
-        )
-
-    mz = materializer.MaterializerHandler(
-        driver=maindb_driver,
-        storage=storage,
-        pubsub=pubsub,
-        check_delay=0.05,
-    )
-    await mz.initialize()
-
-    async with datamanagers.with_transaction() as txn:
-        assert (
-            await datamanagers.resources.get_number_of_resources(
-                txn, kbid=knowledgebox_ingest
-            )
-            == -1
-        )
-        assert (
-            await datamanagers.resources.calculate_number_of_resources(
-                txn, kbid=knowledgebox_ingest
-            )
-            == count
-        )
-
-    await pubsub.publish(
-        const.PubSubChannels.RESOURCE_NOTIFY.format(kbid=knowledgebox_ingest),
-        writer_pb2.Notification(
-            kbid=knowledgebox_ingest,
-            action=writer_pb2.Notification.Action.COMMIT,
-        ).SerializeToString(),
-    )
-
-    await asyncio.sleep(0.2)
-
-    async with datamanagers.with_transaction() as txn:
-        assert (
-            await datamanagers.resources.get_number_of_resources(
-                txn, kbid=knowledgebox_ingest
-            )
-            == count
-        )
-
-    await mz.finalize()
-
-    assert audit_storage.js.publish.call_count == 1
-    assert audit_storage.js.publish.call_args[0][0] == "kb-usage.nuclia_db"
-    pb = KbUsage()
-    pb.ParseFromString(audit_storage.js.publish.call_args[0][1])
-    assert pb.storage.resources == count
-    assert pb.service == Service.NUCLIA_DB
-    assert pb.kb_id == knowledgebox_ingest

nucliadb/ingest/tests/integration/consumer/test_pull.py

@@ -1,144 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-#
-import asyncio
-import base64
-import time
-import uuid
-from dataclasses import dataclass
-from unittest.mock import patch
-
-import pytest
-from fastapi import FastAPI
-from nucliadb_protos.writer_pb2 import BrokerMessage
-from uvicorn.config import Config  # type: ignore
-from uvicorn.server import Server  # type: ignore
-
-from nucliadb.ingest.consumer.pull import PullWorker
-from nucliadb_utils import const
-from nucliadb_utils.fastapi.run import start_server
-from nucliadb_utils.nats import NatsConnectionManager
-from nucliadb_utils.tests import free_port
-
-pytestmark = pytest.mark.asyncio
-
-
-def create_broker_message(kbid: str) -> BrokerMessage:
-    bm = BrokerMessage()
-    bm.uuid = uuid.uuid4().hex
-    bm.kbid = kbid
-    bm.texts["text1"].body = "My text1"
-    bm.basic.title = "My Title"
-    bm.source == BrokerMessage.MessageSource.PROCESSOR
-
-    return bm
-
-
-@dataclass
-class PullProcessorAPI:
-    url: str
-    messages: list[BrokerMessage]
-
-
-@pytest.fixture()
-async def pull_processor_api():
-    app = FastAPI()
-    messages: list[BrokerMessage] = []  # type: ignore
-
-    @app.get("/api/v1/internal/processing/pull")
-    async def pull():
-        if len(messages) == 0:
-            return {"status": "empty"}
-        message = messages.pop()
-        return {
-            "status": "ok",
-            "payload": base64.b64encode(message.SerializeToString()).decode(),
-            "msgid": str(len(messages)),
-        }
-
-    port = free_port()
-    config = Config(app, host="0.0.0.0", port=port, http="auto")
-    server = Server(config=config)
-
-    await start_server(server, config)
-
-    url = f"http://127.0.0.1:{port}"
-    with patch(
-        "nucliadb.common.http_clients.processing.nuclia_settings.nuclia_processing_cluster_url",
-        url,
-    ):
-        yield PullProcessorAPI(url=url, messages=messages)
-
-    await server.shutdown()
-
-
-@pytest.fixture()
-async def pull_worker(maindb_driver, pull_processor_api: PullProcessorAPI):
-    worker = PullWorker(
-        driver=maindb_driver,
-        partition="1",
-        storage=None,  # type: ignore
-        pull_time_error_backoff=5,
-        pull_time_empty_backoff=0.1,
-    )
-
-    task = asyncio.create_task(worker.loop())
-    yield worker
-    task.cancel()
-
-
-async def wait_for_messages(messages: list[BrokerMessage], max_time: int = 10) -> None:
-    start = time.monotonic()
-    while time.monotonic() - start < max_time:
-        if len(messages) == 0:
-            await asyncio.sleep(
-                0.1
-            )  # extra sleep to make sure it's flushed to consumer
-            return
-
-        await asyncio.sleep(0.1)
-
-
-async def test_pull_full_integration(
-    ingest_consumers,
-    ingest_processed_consumer,
-    pull_worker: PullWorker,
-    pull_processor_api: PullProcessorAPI,
-    knowledgebox_ingest: str,
-    nats_manager: NatsConnectionManager,
-):
-    # make sure stream is empty
-    consumer_info1 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST.name, const.Streams.INGEST.group.format(partition="1")
-    )
-    consumer_info2 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST_PROCESSED.name, const.Streams.INGEST_PROCESSED.group
-    )
-    assert consumer_info1.delivered.stream_seq == 0
-    assert consumer_info2.delivered.stream_seq == 0
-
-    # add message that should go to first consumer
-    pull_processor_api.messages.append(create_broker_message(knowledgebox_ingest))
-    await wait_for_messages(pull_processor_api.messages)
-
-    consumer_info1 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST.name, const.Streams.INGEST_PROCESSED.group
-    )
-
-    assert consumer_info1.delivered.stream_seq == 1

nucliadb/ingest/tests/integration/consumer/test_service.py

@@ -1,81 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-#
-import uuid
-
-import pytest
-from nucliadb_protos.writer_pb2 import BrokerMessage
-
-from nucliadb_utils import const
-from nucliadb_utils.nats import NatsConnectionManager
-from nucliadb_utils.transaction import TransactionUtility
-
-pytestmark = pytest.mark.asyncio
-
-
-def create_broker_message(kbid: str) -> BrokerMessage:
-    bm = BrokerMessage()
-    bm.uuid = uuid.uuid4().hex
-    bm.kbid = kbid
-    bm.texts["text1"].body = "My text1"
-    bm.basic.title = "My Title"
-
-    return bm
-
-
-async def test_separated_ingest_consumer(
-    ingest_consumers,
-    ingest_processed_consumer,
-    knowledgebox_ingest,
-    transaction_utility: TransactionUtility,
-    nats_manager: NatsConnectionManager,
-    pubsub,
-):
-    bm_normal = create_broker_message(knowledgebox_ingest)
-    bm_processed = create_broker_message(knowledgebox_ingest)
-    bm_processed.source == BrokerMessage.MessageSource.PROCESSOR
-
-    await transaction_utility.commit(bm_normal, partition=1, wait=True)
-
-    consumer_info1 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST.name, const.Streams.INGEST.group.format(partition="1")
-    )
-    consumer_info2 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST_PROCESSED.name, const.Streams.INGEST_PROCESSED.group
-    )
-
-    assert consumer_info1.delivered.stream_seq == 1
-    assert consumer_info2.delivered.stream_seq == 0
-
-    await transaction_utility.commit(
-        bm_normal,
-        partition=1,
-        wait=True,
-        target_subject=const.Streams.INGEST_PROCESSED.subject,
-    )
-
-    consumer_info1 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST.name, const.Streams.INGEST.group.format(partition="1")
-    )
-    consumer_info2 = await nats_manager.js.consumer_info(
-        const.Streams.INGEST_PROCESSED.name, const.Streams.INGEST_PROCESSED.group
-    )
-
-    assert consumer_info1.delivered.stream_seq == 1
-    assert consumer_info2.delivered.stream_seq == 2

nucliadb/ingest/tests/integration/consumer/test_shard_creator.py

@@ -1,68 +0,0 @@
-# Copyright (C) 2021 Bosutech XXI S.L.
-#
-# nucliadb is offered under the AGPL v3.0 and as commercial software.
-# For commercial licensing, contact us at info@nuclia.com.
-#
-# AGPL:
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of the
-# License, or (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
-# GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU Affero General Public License
-# along with this program. If not, see <http://www.gnu.org/licenses/>.
-#
-
-import asyncio
-
-import pytest
-
-from nucliadb.ingest.consumer import shard_creator
-from nucliadb_protos import writer_pb2
-from nucliadb_utils import const
-
-pytestmark = pytest.mark.asyncio
-
-
-async def test_shard_auto_create(
-    maindb_driver,
-    pubsub,
-    storage,
-    fake_node,
-    knowledgebox_ingest,
-):
-    from nucliadb.common.cluster.settings import settings
-
-    settings.max_shard_paragraphs = 1
-
-    sc = shard_creator.ShardCreatorHandler(
-        driver=maindb_driver,
-        storage=storage,
-        pubsub=pubsub,
-        check_delay=0.05,
-    )
-    await sc.initialize()
-
-    original_kb_shards = await sc.shard_manager.get_shards_by_kbid_inner(
-        knowledgebox_ingest
-    )
-
-    await pubsub.publish(
-        const.PubSubChannels.RESOURCE_NOTIFY.format(kbid=knowledgebox_ingest),
-        writer_pb2.Notification(
-            kbid=knowledgebox_ingest,
-            action=writer_pb2.Notification.Action.INDEXED,
-        ).SerializeToString(),
-    )
-
-    await asyncio.sleep(0.2)
-
-    await sc.finalize()
-
-    kb_shards = await sc.shard_manager.get_shards_by_kbid_inner(knowledgebox_ingest)
-    assert len(kb_shards.shards) == len(original_kb_shards.shards) + 1