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

nucliadb/tests/unit/export_import/test_utils.py

@@ -1,301 +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 io import BytesIO
-from unittest.mock import AsyncMock, Mock
-
-import nats.errors
-import pytest
-from nucliadb_protos.writer_pb2 import BrokerMessage, BrokerMessageBlobReference
-from starlette.requests import Request
-
-from nucliadb.export_import.exceptions import ExportStreamExhausted
-from nucliadb.export_import.models import ImportMetadata
-from nucliadb.export_import.utils import (
-    ExportStream,
-    TaskRetryHandler,
-    get_cloud_files,
-    import_broker_message,
-    transaction_commit,
-)
-from nucliadb_models.export_import import Status
-from nucliadb_protos import resources_pb2
-from nucliadb_utils.const import Streams
-
-
-@pytest.fixture(scope="function")
-def transaction():
-    mock = Mock()
-    mock.commit = AsyncMock()
-    yield mock
-
-
-@pytest.fixture(scope="function")
-def partitioning():
-    mock = Mock()
-    mock.generate_partition = Mock(return_value=1)
-    yield mock
-
-
-def get_cf(uri=None) -> resources_pb2.CloudFile:
-    cf = resources_pb2.CloudFile()
-    uri = uri or "//foo/bar"
-    cf.uri = uri
-    cf.source = resources_pb2.CloudFile.Source.LOCAL
-    return cf
-
-
-@pytest.fixture(scope="function")
-def broker_message():
-    bm = BrokerMessage()
-    bm.kbid = "foobar"
-
-    # Add a file field
-    file = resources_pb2.FieldFile()
-    file.file.CopyFrom(get_cf("file"))
-    bm.files["file"].CopyFrom(file)
-
-    # Add a conversation with an attachment
-    conversation = resources_pb2.Conversation()
-    message = resources_pb2.Message()
-    attachment = get_cf("attachment")
-    message.content.attachments.append(attachment)
-    conversation.messages.append(message)
-    bm.conversations["conversation"].CopyFrom(conversation)
-
-    # Add a layout with a file
-    layout = resources_pb2.FieldLayout()
-    block = resources_pb2.Block()
-    block.file.CopyFrom(get_cf("layout"))
-    layout.body.blocks["foo"].CopyFrom(block)
-    bm.layouts["layout"].CopyFrom(layout)
-
-    # Field extracted data
-    fed = resources_pb2.FileExtractedData()
-    fed.file_generated["foo"].CopyFrom(get_cf("field_file_generated"))
-    fed.file_preview.CopyFrom(get_cf("field_file_preview"))
-    fed.file_thumbnail.CopyFrom(get_cf("field_file_thumbnail"))
-    fed.file_pages_previews.pages.append(get_cf("field_file_pages_previews"))
-    bm.file_extracted_data.append(fed)
-
-    # Link extracted data
-    led = resources_pb2.LinkExtractedData()
-    led.link_thumbnail.CopyFrom(get_cf("link_thumbnail"))
-    led.link_preview.CopyFrom(get_cf("link_preview"))
-    led.link_image.CopyFrom(get_cf("link_image"))
-    bm.link_extracted_data.append(led)
-
-    # Field metadata
-    fcmw = resources_pb2.FieldComputedMetadataWrapper()
-    fcmw.metadata.metadata.thumbnail.CopyFrom(get_cf("metadata_thumbnail"))
-    fcmw.metadata.split_metadata["foo"].thumbnail.CopyFrom(
-        get_cf("metadata_split_thumbnail")
-    )
-    bm.field_metadata.append(fcmw)
-
-    return bm
-
-
-class ContextMock:
-    def __init__(self, transaction, partitioning):
-        self.transaction = transaction
-        self.partitioning = partitioning
-
-
-async def test_import_broker_message(broker_message, transaction, partitioning):
-    context = ContextMock(transaction, partitioning)
-
-    import_kbid = "import_kbid"
-    assert broker_message.kbid != import_kbid
-
-    await import_broker_message(context, import_kbid, broker_message)
-
-    # Sends two messages
-    assert transaction.commit.call_count == 2
-
-    for call in transaction.commit.call_args_list:
-        # Message contains import kbid
-        assert call[0][0].kbid == import_kbid
-
-        # Sends to correct topic
-        assert call[1]["target_subject"] == Streams.INGEST_PROCESSED.subject
-
-
-def test_get_cloud_files(broker_message):
-    # All expected binaries are returned
-    binaries = get_cloud_files(broker_message)
-    assert len(binaries) == 12
-    for cf in binaries:
-        assert cf.source == resources_pb2.CloudFile.Source.LOCAL
-
-    # Make sure that the source is set to export on the broker message cfs
-    for cf in get_cloud_files(broker_message):
-        assert cf.source == resources_pb2.CloudFile.Source.EXPORT
-
-
-async def test_export_stream_simple():
-
-    async def export_generator():
-        export = BytesIO(b"1234567890")
-        while True:
-            await asyncio.sleep(0)
-            chunk = export.read(2)
-            if not chunk:
-                break
-            yield chunk
-
-    stream = ExportStream(export_generator())
-    assert stream.read_bytes == 0
-    assert await stream.read(0) == b""
-    assert stream.read_bytes == 0
-    assert await stream.read(1) == b"1"
-    assert stream.read_bytes == 1
-    assert await stream.read(2) == b"23"
-    assert stream.read_bytes == 3
-    assert await stream.read(50) == b"4567890"
-    assert stream.read_bytes == 10
-    with pytest.raises(ExportStreamExhausted):
-        await stream.read(1)
-
-
-class DummyTestRequest(Request):
-    def __init__(self, data: bytes, receive_chunk_size: int = 10):
-        super().__init__(
-            scope={
-                "type": "http",
-                "http_version": "1.1",
-                "method": "GET",
-                "headers": [],
-            },
-            receive=self.receive,
-        )
-        self.receive_chunk_size = receive_chunk_size
-        self.bytes = BytesIO(data)
-
-    async def receive(self):
-        chunk = self.bytes.read(self.receive_chunk_size)
-        more_data = True
-        if chunk == b"":
-            more_data = False
-        return {"type": "http.request", "body": chunk, "more_body": more_data}
-
-
-async def test_export_stream():
-    request = DummyTestRequest(data=b"01234XYZ", receive_chunk_size=2)
-
-    export_stream = ExportStream(request.stream())
-    assert await export_stream.read(0) == b""
-    assert export_stream.read_bytes == 0
-
-    for i in range(5):
-        assert await export_stream.read(1) == f"{i}".encode()
-    assert export_stream.read_bytes == 5
-
-    assert await export_stream.read(3) == b"XYZ"
-    assert export_stream.read_bytes == 8
-
-    with pytest.raises(ExportStreamExhausted):
-        await export_stream.read(1)
-
-    with pytest.raises(ExportStreamExhausted):
-        await export_stream.read(0)
-
-    request = DummyTestRequest(data=b"foobar", receive_chunk_size=2)
-    export_stream = ExportStream(request.stream())
-    assert await export_stream.read(50) == b"foobar"
-    assert export_stream.read_bytes == 6
-
-    with pytest.raises(ExportStreamExhausted):
-        await export_stream.read(0)
-
-
-class TestTaskRetryHandler:
-    @pytest.fixture(scope="function")
-    def callback(self):
-        return AsyncMock()
-
-    @pytest.fixture(scope="function")
-    def dm(self):
-        dm = Mock()
-        dm.set_metadata = AsyncMock()
-        return dm
-
-    @pytest.fixture(scope="function")
-    def metadata(self):
-        return ImportMetadata(kbid="kbid", id="import_id")
-
-    async def test_ok(self, callback, dm, metadata):
-        callback.return_value = 100
-        trh = TaskRetryHandler("foo", dm, metadata)
-        callback_retried = trh.wrap(callback)
-
-        result = await callback_retried("foo", bar="baz")
-        assert result == 100
-
-        callback.assert_called_once_with("foo", bar="baz")
-
-        assert metadata.task.status == Status.FINISHED
-
-    async def test_errors_are_retried(self, callback, dm, metadata):
-        callback.side_effect = ValueError("foo")
-
-        trh = TaskRetryHandler("foo", dm, metadata, max_tries=2)
-        callback_retried = trh.wrap(callback)
-
-        with pytest.raises(ValueError):
-            await callback_retried("foo", bar="baz")
-
-        callback.assert_called_once_with("foo", bar="baz")
-
-        assert metadata.task.status == Status.RUNNING
-        assert metadata.task.retries == 1
-
-        with pytest.raises(ValueError):
-            await callback_retried("foo", bar="baz")
-
-        assert metadata.task.status == Status.RUNNING
-        assert metadata.task.retries == 2
-
-    async def test_ignored_statuses(self, callback, dm, metadata):
-        trh = TaskRetryHandler("foo", dm, metadata)
-        callback_retried = trh.wrap(callback)
-
-        for status in (Status.ERRORED, Status.FINISHED):
-            metadata.task.status = status
-            await callback_retried("foo", bar="baz")
-            callback.assert_not_called()
-
-
-async def test_transaction_commit_sends_storage_reference_on_max_payload_error():
-    context = Mock()
-    context.transaction.commit = AsyncMock(
-        side_effect=[nats.errors.MaxPayloadError, None]
-    )
-    context.blob_storage = AsyncMock()
-    context.blob_storage.set_stream_message.return_value = "key"
-
-    bm = BrokerMessage(kbid="kbid", uuid="uuid")
-
-    await transaction_commit(context, bm, 1)
-
-    assert context.transaction.commit.call_count == 2
-    call = context.transaction.commit.call_args_list[-1]
-    assert isinstance(call[1]["writer"], BrokerMessageBlobReference)
-    assert call[1]["headers"] == {"X-MESSAGE-TYPE": "PROXY"}

nucliadb/tests/unit/migrator/__init__.py

@@ -1,19 +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/tests/unit/migrator/test_migrator.py

@@ -1,87 +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 types
-from unittest.mock import AsyncMock, Mock, patch
-
-import pytest
-
-from nucliadb.migrator import command, migrator
-from nucliadb.migrator.exceptions import MigrationValidationError
-from nucliadb.migrator.models import Migration
-
-
-def test_get_migrations():
-    migrations = migrator.get_migrations()
-    assert len(migrations) > 0
-    assert migrations[0].version == 1
-    assert migrations[0].module.__name__ == "migrations.0001_bootstrap"
-
-
-def test_get_migration_with_filtering():
-    with patch("nucliadb.migrator.utils.get_migration_modules") as mock:
-        mock.return_value = [
-            (types.ModuleType("m1"), 1),
-            (types.ModuleType("m2"), 2),
-            (types.ModuleType("m3"), 3),
-            (types.ModuleType("m4"), 4),
-        ]
-        migrations = migrator.get_migrations(from_version=2, to_version=3)
-        assert len(migrations) == 1
-        assert migrations[0].version == 3
-        assert migrations[0].module.__name__ == "m3"
-
-
-async def test_run_all_kb_migrations_raises_on_failure():
-    execution_context = Mock()
-    execution_context.data_manager = Mock()
-    execution_context.data_manager.get_kb_migrations = AsyncMock(
-        return_value=["foo", "bar"]
-    )
-    execution_context.settings = Mock(max_concurrent_migrations=1)
-    with patch(
-        "nucliadb.migrator.migrator.run_kb_migrations",
-        side_effect=[None, Exception("Boom")],
-    ) as mock:
-        with pytest.raises(Exception) as exc_info:
-            await migrator.run_all_kb_migrations(execution_context, 1)
-        assert "Failed to migrate KBs. Failures: 1" in str(exc_info.value)
-        assert mock.call_count == 2
-
-
-async def test_migrations_validation():
-    migrations = [
-        Migration(version=1, module=Mock()),
-        Migration(version=2, module=Mock()),
-        Migration(version=3, module=Mock()),
-    ]
-    with patch("nucliadb.migrator.command.get_migrations", return_value=migrations):
-        command.validate()
-
-
-async def test_migrations_validation_with_errors():
-    migrations = [
-        Migration(version=1, module=Mock()),
-        Migration(version=2, module=Mock()),
-        Migration(version=2, module=Mock()),
-        Migration(version=3, module=Mock()),
-    ]
-    with patch("nucliadb.migrator.command.get_migrations", return_value=migrations):
-        with pytest.raises(MigrationValidationError):
-            command.validate()

nucliadb/tests/unit/tasks/__init__.py

@@ -1,19 +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/tests/unit/tasks/conftest.py

@@ -1,42 +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/>.
-#
-from unittest.mock import AsyncMock, Mock
-
-import nats
-import pytest
-
-
-@pytest.fixture(scope="function")
-def nats_manager():
-    nats_manager = Mock()
-    nats_manager.subscribe = AsyncMock()
-    js = Mock()
-    js.stream_info = AsyncMock(side_effect=nats.js.errors.NotFoundError)
-    js.add_stream = AsyncMock()
-    nats_manager.js = js
-    yield nats_manager
-
-
-@pytest.fixture(scope="function")
-def context(nats_manager):
-    context = Mock()
-    context.initialize = AsyncMock()
-    context.nats_manager = nats_manager
-    yield context

nucliadb/tests/unit/tasks/test_consumer.py

@@ -1,92 +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/>.
-#
-from unittest.mock import AsyncMock, MagicMock
-
-import pydantic
-import pytest
-
-from nucliadb.tasks.consumer import create_consumer
-
-
-class Message(pydantic.BaseModel):
-    kbid: str
-
-
-def test_create_consumer():
-    stream = MagicMock()
-
-    async def callback(): ...
-
-    consumer = create_consumer(
-        "foo", stream=stream, msg_type=Message, callback=callback
-    )
-    assert not consumer.initialized
-
-    assert consumer.name == "foo"
-    assert consumer.stream == stream
-    assert consumer.callback == callback
-    assert consumer.msg_type == Message
-
-
-class TestSubscriptionWorker:
-    @pytest.fixture(scope="function")
-    async def callback(self):
-        yield AsyncMock()
-
-    @pytest.fixture(scope="function")
-    async def consumer(self, context, callback):
-        consumer = create_consumer(
-            "foo", stream=MagicMock(), callback=callback, msg_type=Message
-        )
-        await consumer.initialize(context)
-        yield consumer
-
-    @pytest.fixture(scope="function")
-    def task_message(self):
-        yield Message(kbid="kbid")
-
-    @pytest.fixture(scope="function")
-    def msg(self, task_message):
-        data = task_message.json().encode("utf-8")
-        msg = MagicMock()
-        msg.data = data
-        msg.ack = AsyncMock()
-        msg.nak = AsyncMock()
-        yield msg
-
-    async def test_callback_ok(self, consumer, msg, callback):
-        await consumer.subscription_worker(msg)
-
-        callback.assert_called_once()
-
-    async def test_callback_error(self, consumer, msg, callback):
-        callback.side_effect = Exception("foo")
-
-        await consumer.subscription_worker(msg)
-
-        callback.assert_called_once()
-        msg.nak.assert_called_once()
-
-    async def test_validation_error(self, consumer, msg):
-        msg.data = b"foo"
-
-        await consumer.subscription_worker(msg)
-
-        msg.ack.assert_called_once()

nucliadb/tests/unit/tasks/test_producer.py

@@ -1,95 +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/>.
-#
-from unittest.mock import AsyncMock, MagicMock, Mock
-
-import nats
-import pydantic
-import pytest
-
-from nucliadb.tasks.producer import create_producer
-
-
-class Message(pydantic.BaseModel):
-    kbid: str
-
-
-def test_create_producer():
-    stream = MagicMock()
-
-    producer = create_producer("foo", stream=stream, msg_type=Message)
-    assert not producer.initialized
-
-    assert producer.name == "foo"
-    assert producer.stream == stream
-
-
-class TestProducer:
-    @pytest.fixture(scope="function")
-    def stream(self):
-        return MagicMock()
-
-    @pytest.fixture(scope="function")
-    def nats_manager(self):
-        mgr = MagicMock()
-        mgr.js.stream_info = AsyncMock(side_effect=nats.js.errors.NotFoundError)
-        mgr.js.add_stream = AsyncMock()
-        mgr.js.publish = AsyncMock()
-        yield mgr
-
-    @pytest.fixture(scope="function")
-    async def producer(self, context, stream, nats_manager):
-        async def callback(context, msg: Message):
-            pass
-
-        producer = create_producer("foo", stream=stream, msg_type=Message)
-        await producer.initialize(context)
-        producer.context.nats_manager = nats_manager
-        yield producer
-
-    async def test_initialize_creates_stream(self, producer, nats_manager):
-        # Check that the stream is on inialization
-        assert nats_manager.js.add_stream.call_count == 1
-        assert nats_manager.js.add_stream.call_args[1]["name"] == producer.stream.name
-        assert nats_manager.js.add_stream.call_args[1]["subjects"] == [
-            producer.stream.subject
-        ]
-
-    async def test_produce_raises_error_if_not_initialized(self, producer):
-        producer.initialized = False
-        with pytest.raises(RuntimeError):
-            await producer(Mock())
-
-    async def test_produce_ok(self, producer, stream):
-        msg = Message(kbid="kbid")
-
-        await producer(msg)
-
-        publish_args = producer.context.nats_manager.js.publish.call_args[0]
-        assert publish_args[0] == stream.subject
-
-        raw_message = publish_args[1]
-        sent_message = Message.parse_raw(raw_message)
-        assert sent_message == msg
-
-    async def test_produce_raises_publish_errors(self, producer, nats_manager):
-        nats_manager.js.publish.side_effect = ValueError("foo")
-
-        with pytest.raises(ValueError):
-            await producer(Mock())