nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/tests/unit/common/cluster/test_rollover.py

@@ -1,279 +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 datetime import datetime
-from unittest.mock import ANY, AsyncMock, MagicMock, patch
-
-import pytest
-
-from nucliadb.common.cluster import manager, rollover
-from nucliadb.common.cluster.index_node import IndexNode
-from nucliadb_protos import writer_pb2
-
-pytestmark = pytest.mark.asyncio
-
-
-@pytest.fixture()
-def available_nodes():
-    nodes = {
-        "0": IndexNode(
-            id="0", address="node-0", available_disk=100, shard_count=0, dummy=True
-        ),
-        "1": IndexNode(
-            id="1", address="node-1", available_disk=100, shard_count=0, dummy=True
-        ),
-        "2": IndexNode(
-            id="2", address="node-2", available_disk=100, shard_count=0, dummy=True
-        ),
-    }
-    with patch.object(manager, "INDEX_NODES", new=nodes):
-        yield nodes
-
-
-@pytest.fixture()
-def shards():
-    yield writer_pb2.Shards(
-        shards=[
-            writer_pb2.ShardObject(
-                shard="1",
-                replicas=[
-                    writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="1")),
-                    writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="2")),
-                ],
-                read_only=True,
-            ),
-            writer_pb2.ShardObject(
-                shard="2",
-                replicas=[
-                    writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="3")),
-                    writer_pb2.ShardReplica(shard=writer_pb2.ShardCreated(id="4")),
-                ],
-                read_only=False,
-            ),
-        ],
-        kbid="kbid",
-        actual=1,
-    )
-
-
-@pytest.fixture()
-def resource_ids():
-    yield ["1", "2", "3"]
-
-
-@pytest.fixture()
-def resources_datamanager(resource_ids):
-    mock = MagicMock()
-
-    async def iterate_resource_ids(kbid):
-        for id in resource_ids:
-            yield id
-
-    mock.iterate_resource_ids = iterate_resource_ids
-    mock.get_resource_shard_id = AsyncMock()
-    mock.get_resource_shard_id.return_value = "1"
-
-    res = MagicMock()
-
-    res.basic.modified.ToDatetime.return_value = datetime.now()
-
-    mock.get_resource = AsyncMock()
-    mock.get_resource.return_value = res
-
-    mock.get_resource_index_message = AsyncMock()
-    metadata = MagicMock()
-    metadata.modified.ToDatetime.return_value = datetime.now()
-    mock.get_resource_index_message.return_value = metadata
-
-    with patch("nucliadb.common.cluster.rollover.datamanagers.resources", mock):
-        yield mock
-
-
-@pytest.fixture()
-def cluster_datamanager(resource_ids, shards):
-    mock = MagicMock()
-    mock.get_kb_shards = AsyncMock()
-    mock.get_kb_shards.return_value = shards
-    mock.update_kb_shards = AsyncMock()
-
-    with patch("nucliadb.common.cluster.rollover.datamanagers.cluster", mock):
-        yield mock
-
-
-@pytest.fixture()
-def rollover_datamanager(resource_ids, cluster_datamanager):
-    mock = MagicMock()
-    mock.get_kb_rollover_shards = AsyncMock()
-    mock.get_kb_rollover_shards.return_value = None
-    mock.update_kb_rollover_shards = AsyncMock()
-    mock.delete_kb_rollover_shard = AsyncMock()
-    mock.delete_kb_rollover_shards = AsyncMock()
-    mock.get_to_index = AsyncMock(side_effect=["1", None])
-    mock.add_indexed = AsyncMock()
-    mock.remove_to_index = AsyncMock()
-    mock.get_indexed_data = AsyncMock(return_value=("1", 1))
-    mock.remove_indexed = AsyncMock()
-
-    async def _mock_indexed_keys(kbid):
-        yield "1"
-
-    mock.iter_indexed_keys = _mock_indexed_keys
-
-    with patch("nucliadb.common.cluster.rollover.datamanagers.rollover", mock), patch(
-        "nucliadb.common.cluster.rollover.datamanagers.with_transaction",
-        return_value=AsyncMock(),
-    ), patch(
-        "nucliadb.ingest.consumer.shard_creator.locking.distributed_lock",
-        return_value=AsyncMock(),
-    ):
-        yield mock
-
-
-@pytest.fixture()
-def app_context(rollover_datamanager, resources_datamanager, available_nodes):
-    mock = MagicMock()
-    mock.shard_manager = MagicMock()
-    mock.shard_manager.rollback_shard = AsyncMock()
-    mock.shard_manager.add_resource = AsyncMock()
-    mock.shard_manager.delete_resource = AsyncMock()
-    mock.kv_driver = MagicMock()
-
-    consumer_info = MagicMock()
-    consumer_info.delivered.stream_seq = 0
-    mock.nats_manager.js.consumer_info.return_value = consumer_info
-    yield mock
-
-
-async def test_create_rollover_shards(
-    app_context, available_nodes, shards: writer_pb2.Shards, rollover_datamanager
-):
-    new_shards = await rollover.create_rollover_shards(app_context, "kbid")
-
-    assert new_shards.kbid == "kbid"
-    assert sum(
-        [len(node.writer.calls["NewShard"]) for node in available_nodes.values()]
-    ) == sum([len(s.replicas) for s in shards.shards])
-    rollover_datamanager.update_kb_rollover_shards.assert_called_with(
-        ANY, kbid="kbid", kb_shards=new_shards
-    )
-
-
-async def test_create_rollover_shards_does_not_recreate(
-    app_context, shards: writer_pb2.Shards, rollover_datamanager
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-
-    await rollover.create_rollover_shards(app_context, "kbid")
-
-    app_context.shard_manager.rollback_shard.assert_not_called()
-    rollover_datamanager.update_kb_rollover_shards.assert_not_called()
-
-
-async def test_index_rollover_shards(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-
-    await rollover.index_rollover_shards(app_context, "kbid")
-    rollover_datamanager.add_indexed.assert_called_with(
-        ANY, kbid="kbid", resource_id="1", shard_id="1", modification_time=1
-    )
-
-
-async def test_index_rollover_shards_handles_missing_shards(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = None
-    with pytest.raises(rollover.UnexpectedRolloverError):
-        await rollover.index_rollover_shards(app_context, "kbid")
-
-
-async def test_index_rollover_shards_handles_missing_shard_id(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-    resources_datamanager.get_resource_shard_id.return_value = None
-    with pytest.raises(rollover.UnexpectedRolloverError):
-        await rollover.index_rollover_shards(app_context, "kbid")
-
-
-async def test_index_rollover_shards_handles_missing_res(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-    resources_datamanager.get_resource_index_message.return_value = None
-
-    await rollover.index_rollover_shards(app_context, "kbid")
-
-    rollover_datamanager.remove_to_index.assert_called_with(
-        ANY, kbid="kbid", resource="1"
-    )
-
-
-async def test_cutover_shards(
-    app_context, rollover_datamanager, cluster_datamanager, shards
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-
-    await rollover.cutover_shards(app_context, "kbid")
-
-    cluster_datamanager.update_kb_shards.assert_called_with(
-        ANY, kbid="kbid", shards=ANY
-    )
-    [
-        app_context.shard_manager.rollback_shard.assert_any_call(shard)
-        for shard in shards.shards
-    ]
-
-
-async def test_cutover_shards_missing(app_context, rollover_datamanager):
-    rollover_datamanager.get_kb_rollover_shards.return_value = None
-
-    with pytest.raises(rollover.UnexpectedRolloverError):
-        await rollover.cutover_shards(app_context, "kbid")
-
-
-async def test_validate_indexed_data(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-
-    indexed_res = await rollover.validate_indexed_data(app_context, "kbid")
-    assert len(indexed_res) == len(resource_ids)
-    [
-        resources_datamanager.get_resource_index_message.assert_any_call(
-            ANY, kbid="kbid", rid=res_id
-        )
-        for res_id in resource_ids
-    ]
-
-
-async def test_validate_indexed_data_handles_missing_res(
-    app_context, rollover_datamanager, resources_datamanager, shards, resource_ids
-):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-    resources_datamanager.get_resource_index_message.return_value = None
-    assert len(await rollover.validate_indexed_data(app_context, "kbid")) == 0
-
-
-async def test_rollover_shards(app_context, rollover_datamanager, shards, resource_ids):
-    rollover_datamanager.get_kb_rollover_shards.return_value = shards
-    resource_ids.clear()
-
-    await rollover.rollover_kb_shards(app_context, "kbid")

nucliadb/tests/unit/common/maindb/__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/tests/unit/common/maindb/test_driver.py

@@ -1,127 +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 import mock
-
-import pytest
-
-from nucliadb.common.maindb.driver import Driver
-
-
-class TransactionTest:
-    def __init__(self):
-        self.abort = mock.AsyncMock()
-        self.open = True
-
-    async def commit(self, **kw):
-        self.open = False
-
-
-@pytest.fixture(scope="function")
-def txn():
-    return TransactionTest()
-
-
-@pytest.fixture(scope="function")
-def driver(txn) -> Driver:  # type: ignore
-    driver = Driver()
-    with mock.patch.object(driver, "begin", new=mock.AsyncMock(return_value=txn)):
-        yield driver
-
-
-@pytest.mark.asyncio
-async def test_driver_async_abort(driver, txn):
-    async with driver.transaction(wait_for_abort=False):
-        pass
-
-    assert len(driver._abort_tasks) == 1
-    await asyncio.sleep(0.1)
-
-    txn.abort.assert_called_once()
-    assert len(driver._abort_tasks) == 0
-
-
-@pytest.mark.asyncio
-async def test_driver_finalize_aborts_transactions(driver, txn):
-    async with driver.transaction(wait_for_abort=False):
-        pass
-
-    assert len(driver._abort_tasks) == 1
-
-    await driver.finalize()
-
-    txn.abort.assert_called_once()
-
-    assert len(driver._abort_tasks) == 0
-
-
-@pytest.mark.asyncio
-async def test_transaction_handles_txn_begin_errors(driver):
-    driver.begin.side_effect = ValueError()
-    testmock = mock.AsyncMock()
-    with pytest.raises(ValueError):
-        async with driver.transaction():
-            testmock()
-    testmock.assert_not_called()
-
-
-@pytest.mark.asyncio
-async def test_transaction_does_not_abotr_if_commited(driver):
-    async with driver.transaction() as txn:
-        await txn.commit()
-    txn.abort.assert_not_awaited()
-
-
-@pytest.mark.asyncio
-async def test_transaction_aborts_if_txn_open(driver):
-    async with driver.transaction() as txn:
-        pass
-    txn.abort.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_transaction_aborts_on_errors(driver):
-    with pytest.raises(ValueError):
-        async with driver.transaction() as txn:
-            raise ValueError()
-    txn.abort.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_transaction_wait_for_abort(driver):
-    async with driver.transaction(wait_for_abort=False) as txn:
-        pass
-    txn.abort.assert_not_awaited()
-    txn.abort.assert_called_once()
-    txn.abort.reset_mock()
-
-    # Should wait for abort if there are errors
-    with pytest.raises(ValueError):
-        async with driver.transaction(wait_for_abort=False) as txn:
-            raise ValueError()
-    txn.abort.assert_awaited_once()
-    txn.abort.assert_called_once()
-    txn.abort.reset_mock()
-
-    async with driver.transaction(wait_for_abort=True) as txn:
-        pass
-    txn.abort.assert_awaited_once()
-    txn.abort.assert_called_once()
-    txn.abort.reset_mock()

nucliadb/tests/unit/common/maindb/test_tikv.py

@@ -1,53 +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 pytest
-
-from nucliadb.common.maindb.exceptions import ConflictError
-from nucliadb.common.maindb.tikv import LeaderNotFoundError, TiKVDataLayer
-
-
-@pytest.mark.parametrize(
-    "tikv_exception,handled_exception",
-    [
-        (
-            Exception("gRPC error: RpcFailure: 4-DEADLINE_EXCEEDED Deadline Exceeded"),
-            TimeoutError,
-        ),
-        (Exception("Leader of region 34234 is not found"), LeaderNotFoundError),
-    ],
-)
-async def test_get_retrials(tikv_exception, handled_exception):
-    inner_txn = MagicMock(get=AsyncMock(side_effect=tikv_exception))
-    tikv_txn = TiKVDataLayer(inner_txn)
-
-    with pytest.raises(handled_exception):
-        await tikv_txn.get("key")
-
-    assert inner_txn.get.call_count == 2
-
-
-async def test_commit_raises_conflict_error():
-    inner_txn = MagicMock(commit=AsyncMock(side_effect=Exception("WriteConflict")))
-    tikv_txn = TiKVDataLayer(inner_txn)
-
-    with pytest.raises(ConflictError):
-        await tikv_txn.commit()

nucliadb/tests/unit/common/maindb/test_utils.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/>.
-
-from unittest.mock import AsyncMock, patch
-
-import pytest
-
-from nucliadb.common.maindb.utils import settings, setup_driver
-from nucliadb_utils.exceptions import ConfigurationError
-from nucliadb_utils.utilities import Utility, clean_utility
-
-
-@pytest.fixture(autouse=True)
-def reset_driver_utils():
-    clean_utility(Utility.MAINDB_DRIVER)
-    yield
-    clean_utility(Utility.MAINDB_DRIVER)
-
-
-@pytest.mark.asyncio
-async def test_setup_driver_redis():
-    mock = AsyncMock(initialized=False)
-    with patch.object(settings, "driver", "redis"), patch.object(
-        settings, "driver_redis_url", "driver_redis_url"
-    ), patch("nucliadb.common.maindb.utils.RedisDriver", return_value=mock):
-        assert await setup_driver() == mock
-        mock.initialize.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_setup_driver_tikv():
-    mock = AsyncMock(initialized=False)
-    with patch.object(settings, "driver", "tikv"), patch.object(
-        settings, "driver_tikv_url", "driver_tikv_url"
-    ), patch("nucliadb.common.maindb.utils.TiKVDriver", return_value=mock):
-        assert await setup_driver() == mock
-        mock.initialize.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_setup_driver_pg():
-    mock = AsyncMock(initialized=False)
-    with patch.object(settings, "driver", "pg"), patch.object(
-        settings, "driver_pg_url", "driver_pg_url"
-    ), patch("nucliadb.common.maindb.utils.PGDriver", return_value=mock):
-        assert await setup_driver() == mock
-        mock.initialize.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_setup_driver_local():
-    mock = AsyncMock(initialized=False)
-    with patch.object(settings, "driver", "local"), patch.object(
-        settings, "driver_local_url", "driver_local_url"
-    ), patch("nucliadb.common.maindb.utils.LocalDriver", return_value=mock):
-        assert await setup_driver() == mock
-        mock.initialize.assert_awaited_once()
-
-
-@pytest.mark.asyncio
-async def test_setup_driver_error():
-    with patch.object(settings, "driver", "pg"), patch.object(
-        settings, "driver_pg_url", None
-    ), pytest.raises(ConfigurationError):
-        await setup_driver()

nucliadb/tests/unit/common/test_context.py

@@ -1,36 +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
-
-from nucliadb.common.context import ApplicationContext
-
-
-async def test_initialize_happens_only_once():
-    context = ApplicationContext()
-    context._initialize = AsyncMock()  # noqa
-
-    tasks = []
-    for _ in range(10):
-        tasks.append(context.initialize())
-    await asyncio.gather(*tasks)
-
-    context._initialize.assert_awaited_once()
-    context._initialized is True

nucliadb/tests/unit/export_import/__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/export_import/test_datamanager.py

@@ -1,37 +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
-
-from nucliadb.export_import.datamanager import ExportImportDataManager
-
-
-async def test_try_delete_from_storage():
-    driver = Mock()
-    storage = Mock()
-
-    dm = ExportImportDataManager(driver, storage)
-    dm.delete_export = AsyncMock(side_effect=ValueError())  # type: ignore
-    dm.delete_import = AsyncMock(side_effect=ValueError())  # type: ignore
-
-    await dm.try_delete_from_storage("export", "kbid", "foo")
-    await dm.try_delete_from_storage("import", "kbid", "foo")
-
-    dm.delete_export.assert_called()
-    dm.delete_import.assert_called()