nucliadb-utils 4.0.3.post573__py3-none-any.whl → 4.0.3.post575__py3-none-any.whl

nucliadb_utils/tests/unit/storages/test_storage.py

@@ -1,228 +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 math import ceil
-from unittest.mock import AsyncMock, MagicMock
-
-import pytest
-from nucliadb_protos.noderesources_pb2 import Resource as BrainResource
-from nucliadb_protos.noderesources_pb2 import ResourceID
-from nucliadb_protos.nodewriter_pb2 import IndexMessage
-from nucliadb_protos.resources_pb2 import CloudFile
-
-from nucliadb_utils.storages.local import LocalStorageField
-from nucliadb_utils.storages.storage import (
-    ObjectInfo,
-    Storage,
-    StorageField,
-    iter_and_add_size,
-    iter_in_chunk_size,
-)
-
-
-class TestStorageField:
-    @pytest.fixture
-    def storage(self):
-        yield AsyncMock(source=0)
-
-    @pytest.fixture
-    def field(self):
-        yield MagicMock(uri="uri")
-
-    @pytest.fixture
-    def storage_field(self, storage, field):
-        yield LocalStorageField(storage, "bucket", "fullkey", field)
-
-    @pytest.mark.asyncio
-    async def test_delete(self, storage_field: StorageField, storage):
-        await storage_field.delete()
-        storage.delete_upload.assert_called_once_with("uri", "bucket")
-
-
-class StorageTest(Storage):
-    def __init__(self):
-        self.source = 0
-        self.field_klass = lambda: MagicMock()
-        self.deadletter_bucket = "deadletter_bucket"
-        self.indexing_bucket = "indexing_bucket"
-        self.delete_upload = AsyncMock()
-        self.uploadbytes = AsyncMock()
-        self.move = AsyncMock()
-
-    def get_bucket_name(self, kbid):
-        return "bucket"
-
-    async def iterate_objects(self, bucket_name, prefix):
-        yield ObjectInfo(name="uri")
-
-    async def download(self, bucket_name, uri):
-        br = BrainResource(labels=["label"])
-        yield br.SerializeToString()
-
-    async def create_kb(self, kbid):
-        return True
-
-    async def delete_kb(self, kbid):
-        return True
-
-    async def delete_upload(self, uri, bucket):
-        return True
-
-    async def initialize(self) -> None:
-        pass
-
-    async def finalize(self) -> None:
-        pass
-
-    async def schedule_delete_kb(self, kbid: str) -> bool:
-        return True
-
-
-class TestStorage:
-    @pytest.fixture
-    def storage(self):
-        yield StorageTest()
-
-    @pytest.mark.asyncio
-    async def test_delete_resource(self, storage: StorageTest):
-        await storage.delete_resource("bucket", "uri")
-
-        storage.delete_upload.assert_called_once_with("uri", "bucket")
-
-    @pytest.mark.asyncio
-    async def test_indexing(self, storage: StorageTest):
-        msg = BrainResource(resource=ResourceID(uuid="uuid"))
-        await storage.indexing(msg, 1, "1", "kb", "shard")
-
-        storage.uploadbytes.assert_called_once_with(
-            "indexing_bucket", "index/kb/shard/uuid/1", msg.SerializeToString()
-        )
-
-    @pytest.mark.asyncio
-    async def test_reindexing(self, storage: StorageTest):
-        msg = BrainResource(resource=ResourceID(uuid="uuid"))
-        await storage.reindexing(msg, "reindex_id", "1", "kb", "shard")
-
-        storage.uploadbytes.assert_called_once_with(
-            "indexing_bucket", "index/kb/shard/uuid/reindex_id", msg.SerializeToString()
-        )
-
-    @pytest.mark.asyncio
-    async def test_get_indexing(self, storage: StorageTest):
-        im = IndexMessage()
-        im.node = "node"
-        im.shard = "shard"
-        im.txid = 0
-        assert isinstance(await storage.get_indexing(im), BrainResource)
-
-    @pytest.mark.asyncio
-    async def test_get_indexing_storage_key(self, storage: StorageTest):
-        im = IndexMessage()
-        im.node = "node"
-        im.shard = "shard"
-        im.txid = 0
-        im.storage_key = "index/kb/uuid/1"
-        assert isinstance(await storage.get_indexing(im), BrainResource)
-
-    @pytest.mark.asyncio
-    async def test_delete_indexing(self, storage: StorageTest):
-        im = IndexMessage()
-        im.node = "node"
-        im.txid = 0
-        im.storage_key = "index/kb/uuid/1"
-        await storage.delete_indexing(
-            resource_uid="resource_uid", txid=1, kb="kb", logical_shard="logical_shard"
-        )
-
-        storage.uploadbytes.assert_called_once()
-
-    @pytest.mark.asyncio
-    async def test_download_pb(self, storage: StorageTest):
-        assert isinstance(
-            await storage.download_pb(
-                LocalStorageField(storage, "bucket", "fullkey"), BrainResource
-            ),
-            BrainResource,
-        )
-
-    @pytest.mark.asyncio
-    async def test_indexing_bucket_none_attributeerrror(self, storage: StorageTest):
-        storage.indexing_bucket = None
-        msg = BrainResource()
-        im = IndexMessage(node="node", shard="shard", txid=0)
-
-        with pytest.raises(AttributeError):
-            await storage.indexing(msg, 1, "1", "kb", "shard")
-
-        with pytest.raises(AttributeError):
-            await storage.reindexing(msg, "reindex_id", "1", "kb", "shard")
-
-        with pytest.raises(AttributeError):
-            await storage.get_indexing(im)
-
-        with pytest.raises(AttributeError):
-            await storage.delete_indexing(
-                resource_uid="resource_uid",
-                txid=1,
-                kb="kb",
-                logical_shard="logical_shard",
-            )
-
-
-async def testiter_and_add_size():
-    cf = CloudFile()
-
-    async def iter():
-        yield b"foo"
-        yield b"bar"
-
-    cf.size = 0
-    async for _ in iter_and_add_size(iter(), cf):
-        pass
-
-    assert cf.size == 6
-
-
-async def test_iter_in_chunk_size():
-    async def iterable(total_size, *, chunk_size=1):
-        data = b"0" * total_size
-        for i in range(ceil(total_size / chunk_size)):
-            chunk = data[i * chunk_size : (i + 1) * chunk_size]
-            yield chunk
-
-    chunks = [chunk async for chunk in iter_in_chunk_size(iterable(10), chunk_size=4)]
-    assert len(chunks) == 3
-    assert len(chunks[0]) == 4
-    assert len(chunks[1]) == 4
-    assert len(chunks[2]) == 2
-
-    chunks = [chunk async for chunk in iter_in_chunk_size(iterable(0), chunk_size=4)]
-    assert len(chunks) == 0
-
-    # Try with an iterable that yields chunks bigger than the chunk size
-    chunks = [
-        chunk
-        async for chunk in iter_in_chunk_size(
-            iterable(total_size=12, chunk_size=10), chunk_size=4
-        )
-    ]
-    assert len(chunks) == 3
-    assert len(chunks[0]) == 4
-    assert len(chunks[1]) == 4
-    assert len(chunks[2]) == 4

nucliadb_utils/tests/unit/test_asyncio_utils.py

@@ -1,72 +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 random
-import time
-
-import pytest
-
-from nucliadb_utils.asyncio_utils import ConcurrentRunner, run_concurrently
-
-
-async def test_run_concurrently():
-    async def mycoro(value):
-        await asyncio.sleep(random.uniform(0.0, 0.5))
-        return value
-
-    tasks = [mycoro(i) for i in range(10)]
-    results = await run_concurrently(tasks)
-    assert results == list(range(10))
-
-
-async def test_concurrent_runner():
-    async def mycoro(value):
-        await asyncio.sleep(0.5)
-        return value
-
-    # Test without max_tasks: execution time should be less than 1 second
-    # as both should be executed in concurrently
-    runner = ConcurrentRunner()
-    runner.schedule(mycoro(1))
-    runner.schedule(mycoro(2))
-    start = time.perf_counter()
-    results = await runner.wait()
-    end = time.perf_counter()
-    assert results == [1, 2]
-    assert end - start < 1
-
-    # Test with max_tasks: execution time should be at least 1 second
-    runner = ConcurrentRunner(max_tasks=1)
-    runner.schedule(mycoro(1))
-    runner.schedule(mycoro(2))
-    start = time.perf_counter()
-    results = await runner.wait()
-    end = time.perf_counter()
-    assert results == [1, 2]
-    assert end - start >= 1
-
-    # Exception should be raised if any of the coroutines raises an exception
-    with pytest.raises(ValueError):
-
-        async def coro_with_exception():
-            raise ValueError("Test")
-
-        runner = ConcurrentRunner()
-        runner.schedule(coro_with_exception())
-        await runner.wait()

nucliadb_utils/tests/unit/test_authentication.py

@@ -1,159 +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, patch
-
-import pytest
-from fastapi import WebSocket
-from starlette.exceptions import HTTPException
-from starlette.requests import Request
-
-from nucliadb_utils import authentication
-
-
-class TestNucliaCloudAuthenticationBackend:
-    @pytest.fixture()
-    def backend(self):
-        return authentication.NucliaCloudAuthenticationBackend()
-
-    @pytest.fixture()
-    def req(self):
-        return Mock(headers={})
-
-    @pytest.mark.asyncio
-    async def test_authenticate(
-        self, backend: authentication.NucliaCloudAuthenticationBackend, req
-    ):
-        assert await backend.authenticate(req) is None
-
-    @pytest.mark.asyncio
-    async def test_authenticate_with_user(
-        self, backend: authentication.NucliaCloudAuthenticationBackend, req
-    ):
-        req.headers = {
-            backend.roles_header: "admin",
-            backend.user_header: "user",
-        }
-
-        creds, user = await backend.authenticate(req)
-
-        assert creds.scopes == ["admin"]
-        assert user.username == "user"
-
-    @pytest.mark.asyncio
-    async def test_authenticate_with_anon(
-        self, backend: authentication.NucliaCloudAuthenticationBackend, req
-    ):
-        req.headers = {backend.roles_header: "admin"}
-
-        creds, user = await backend.authenticate(req)
-
-        assert creds.scopes == ["admin"]
-        assert user.username == "anonymous"
-
-
-def test_has_required_scope():
-    conn = Mock(auth=Mock(scopes=["admin"]))
-    assert authentication.has_required_scope(conn, ["admin"])
-    assert not authentication.has_required_scope(conn, ["foobar"])
-
-
-class TestRequires:
-    def test_requires_sync(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        assert authentication.requires(["admin"])(lambda request: None)(req) is None
-
-        # test passed as kwargs
-        assert (
-            authentication.requires(["admin"])(lambda request: None)(request=req)
-            is None
-        )
-
-    def test_requires_sync_returns_status(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        with pytest.raises(HTTPException):
-            assert authentication.requires(["foobar"])(lambda request: None)(req)
-
-    def test_requires_sync_returns_redirect(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        with patch.object(req, "url_for", return_value="http://foobar"):
-            resp = authentication.requires(["foobar"], redirect="/foobar")(
-                lambda request: None
-            )(req)
-        assert resp.status_code == 303
-
-    @pytest.mark.asyncio
-    async def test_requires_async(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        async def noop(request): ...
-
-        assert await authentication.requires(["admin"])(noop)(request=req) is None
-
-    @pytest.mark.asyncio
-    async def test_requires_async_returns_status(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        async def noop(request): ...
-
-        with pytest.raises(HTTPException):
-            assert await authentication.requires(["foobar"])(noop)(request=req)
-
-    @pytest.mark.asyncio
-    async def test_requires_async_returns_redirect(self):
-        req = Request({"type": "http", "auth": Mock(scopes=["admin"])})
-
-        async def noop(request): ...
-
-        with patch.object(req, "url_for", return_value="http://foobar"):
-            resp = await authentication.requires(["foobar"], redirect="/foobar")(noop)(
-                request=req
-            )
-        assert resp.status_code == 303
-
-    @pytest.mark.asyncio
-    async def test_requires_ws(self):
-        ws = AsyncMock()
-        req = WebSocket(
-            {"type": "websocket", "auth": Mock(scopes=["admin"]), "websocket": ws},
-            receive=AsyncMock(),
-            send=AsyncMock(),
-        )
-
-        async def noop(websocket): ...
-
-        assert await authentication.requires(["admin"])(noop)(req) is None
-
-    @pytest.mark.asyncio
-    async def test_requires_ws_fail(self):
-        req = WebSocket(
-            {"type": "websocket", "auth": Mock(scopes=["admin"])},
-            receive=AsyncMock(),
-            send=AsyncMock(),
-        )
-
-        async def noop(websocket): ...
-
-        with patch.object(req, "close", return_value=None):
-            assert await authentication.requires(["notallowed"])(noop)(req) is None
-
-            req.close.assert_called_once()

nucliadb_utils/tests/unit/test_helpers.py

@@ -1,45 +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 nucliadb_utils.helpers import async_gen_lookahead
-
-
-async def test_async_gen_lookahead():
-    async def gen(n):
-        for i in range(n):
-            yield f"{i}".encode()
-
-    assert [item async for item in async_gen_lookahead(gen(0))] == []
-    assert [item async for item in async_gen_lookahead(gen(1))] == [(b"0", True)]
-    assert [item async for item in async_gen_lookahead(gen(2))] == [
-        (b"0", False),
-        (b"1", True),
-    ]
-
-
-async def test_async_gen_lookahead_last_chunk_is_empty():
-    async def gen():
-        for chunk in [b"empty", None, b"chunk", b""]:
-            yield chunk
-
-    assert [item async for item in async_gen_lookahead(gen())] == [
-        (b"empty", False),
-        (b"chunk", True),
-    ]

nucliadb_utils/tests/unit/test_nats.py

@@ -1,154 +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 time
-from unittest.mock import AsyncMock, MagicMock, patch
-
-import pytest
-
-from nucliadb_utils import nats
-
-pytestmark = pytest.mark.unit
-
-
-class TestNatsConnectionManager:
-    @pytest.fixture()
-    def nats_conn(self):
-        conn = MagicMock()
-        conn.drain = AsyncMock()
-        conn.close = AsyncMock()
-        with patch("nucliadb_utils.nats.nats.connect", return_value=conn):
-            yield conn
-
-    @pytest.fixture()
-    def js(self):
-        conn = AsyncMock()
-        with patch(
-            "nucliadb_utils.nats.get_traced_jetstream",
-            return_value=conn,
-        ):
-            yield conn
-
-    @pytest.fixture()
-    def manager(self, nats_conn, js):
-        yield nats.NatsConnectionManager(service_name="test", nats_servers=["test"])
-
-    async def test_initialize(self, manager: nats.NatsConnectionManager, nats_conn):
-        await manager.initialize()
-
-        assert manager.nc == nats_conn
-
-    async def test_lifecycle_finalize(
-        self, manager: nats.NatsConnectionManager, nats_conn, js
-    ):
-        await manager.initialize()
-
-        cb = AsyncMock()
-        lost_cb = AsyncMock()
-        await manager.subscribe(
-            subject="subject",
-            queue="queue",
-            stream="stream",
-            cb=cb,
-            subscription_lost_cb=lost_cb,
-            flow_control=True,
-        )
-
-        js.subscribe.assert_called_once_with(
-            subject="subject",
-            queue="queue",
-            stream="stream",
-            cb=cb,
-            manual_ack=True,
-            flow_control=True,
-            config=None,
-        )
-
-        await manager.reconnected_cb()
-        lost_cb.assert_called_once()
-
-        await manager.finalize()
-
-        nats_conn.drain.assert_called_once()
-        nats_conn.close.assert_called_once()
-
-    async def test_healthy(self, manager: nats.NatsConnectionManager):
-        await manager.initialize()
-
-        assert manager.healthy()
-
-        manager._healthy = False
-        assert not manager.healthy()
-
-        manager._healthy = True
-        manager._last_unhealthy = time.monotonic() - 100
-        assert not manager.healthy()
-
-        manager._last_unhealthy = None
-        manager._nc.is_connected = False
-        assert manager.healthy()
-        assert manager._last_unhealthy is not None
-
-    async def test_unsubscribe(
-        self, manager: nats.NatsConnectionManager, nats_conn, js
-    ):
-        await manager.initialize()
-
-        cb = AsyncMock()
-        lost_cb = AsyncMock()
-        sub = await manager.subscribe(
-            subject="subject",
-            queue="queue",
-            stream="stream",
-            cb=cb,
-            subscription_lost_cb=lost_cb,
-            flow_control=True,
-        )
-        assert len(manager._subscriptions) == 1
-
-        await manager.unsubscribe(sub)
-
-        sub.unsubscribe.assert_awaited_once()
-        assert len(manager._subscriptions) == 0
-
-        await manager.finalize()
-
-        nats_conn.drain.assert_called_once()
-        nats_conn.close.assert_called_once()
-
-
-async def test_message_progress_updater():
-    in_progress = AsyncMock()
-    msg = MagicMock(in_progress=in_progress, _ackd=False)
-
-    async with nats.MessageProgressUpdater(msg, 0.05):
-        await asyncio.sleep(0.07)
-
-    in_progress.assert_awaited_once()
-
-
-async def test_message_progress_updater_does_not_update_ack():
-    in_progress = AsyncMock()
-    msg = MagicMock(in_progress=in_progress, _ackd=True)
-
-    async with nats.MessageProgressUpdater(msg, 0.05):
-        await asyncio.sleep(0.07)
-
-    in_progress.assert_not_awaited()