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

nucliadb/ingest/tests/integration/ingest/test_processing_engine.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 uuid import uuid4
-
-import pytest
-
-from nucliadb.tests.utils.aiohttp_session import get_mocked_session
-
-
-@pytest.mark.parametrize("onprem", [True, False])
-@pytest.mark.parametrize(
-    "mock_payload",
-    [
-        {"seqid": 1, "account_seq": 1, "queue": "private"},
-        {"seqid": 1, "account_seq": 1, "queue": "shared"},
-        {"seqid": 1, "account_seq": None, "queue": "private"},
-        {"seqid": 1, "account_seq": None, "queue": "shared"},
-        {"seqid": 1, "queue": "private"},
-        {"seqid": 1, "queue": "shared"},
-    ],
-)
-@pytest.mark.asyncio
-async def test_send_to_process(onprem, mock_payload):
-    """
-    Test that send_to_process does not fail
-    """
-
-    from nucliadb.ingest.processing import ProcessingEngine, PushPayload
-
-    fake_nuclia_proxy_url = "http://fake_proxy"
-    processing_engine = ProcessingEngine(
-        onprem=onprem,
-        nuclia_processing_cluster_url=fake_nuclia_proxy_url,
-        nuclia_public_url=fake_nuclia_proxy_url,
-    )
-    await processing_engine.initialize()
-
-    payload = PushPayload(
-        uuid=str(uuid4()), kbid=str(uuid4()), userid=str(uuid4()), partition=0
-    )
-
-    processing_engine.session = get_mocked_session(
-        "POST", 200, json=mock_payload, context_manager=False
-    )
-    await processing_engine.send_to_process(payload, partition=0)
-
-    await processing_engine.finalize()
-
-
-@pytest.mark.parametrize("onprem", [True, False])
-@pytest.mark.asyncio
-async def test_delete_from_processing(onprem):
-    """
-    Test that send_to_process does not fail
-    """
-
-    from nucliadb.ingest.processing import ProcessingEngine
-
-    fake_nuclia_proxy_url = "http://fake_proxy"
-    processing_engine = ProcessingEngine(
-        onprem=onprem,
-        nuclia_processing_cluster_url=fake_nuclia_proxy_url,
-        nuclia_public_url=fake_nuclia_proxy_url,
-    )
-    await processing_engine.initialize()
-
-    processing_engine.session = get_mocked_session(
-        "POST",
-        200,
-        json={"kbid": "kbid", "resource_id": "resource_id"},
-        context_manager=False,
-    )
-    await processing_engine.delete_from_processing(
-        kbid="kbid", resource_id="resource_id"
-    )
-
-    await processing_engine.finalize()

nucliadb/ingest/tests/integration/ingest/test_relations.py

@@ -1,272 +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.resources_pb2 import (
-    Classification,
-    FieldComputedMetadataWrapper,
-    FieldID,
-    FieldText,
-    FieldType,
-)
-from nucliadb_protos.utils_pb2 import Relation, RelationNode
-from nucliadb_protos.writer_pb2 import BrokerMessage
-
-from nucliadb.ingest import SERVICE_NAME
-from nucliadb_utils.utilities import get_indexing, get_storage
-
-
-@pytest.mark.asyncio
-async def test_ingest_relations_indexing(
-    fake_node, local_files, storage, knowledgebox_ingest, processor
-):
-    rid = str(uuid.uuid4())
-    bm = BrokerMessage(
-        kbid=knowledgebox_ingest, uuid=rid, slug="slug-1", type=BrokerMessage.AUTOCOMMIT
-    )
-
-    e0 = RelationNode(value="E0", ntype=RelationNode.NodeType.ENTITY, subtype="")
-    e1 = RelationNode(
-        value="E1", ntype=RelationNode.NodeType.ENTITY, subtype="Official"
-    )
-    e2 = RelationNode(
-        value="E2", ntype=RelationNode.NodeType.ENTITY, subtype="Propaganda"
-    )
-    r0 = Relation(
-        relation=Relation.RelationType.CHILD, source=e1, to=e2, relation_label="R0"
-    )
-    r1 = Relation(
-        relation=Relation.RelationType.ENTITY, source=e0, to=e2, relation_label="R1"
-    )
-    r2 = Relation(
-        relation=Relation.RelationType.CHILD, source=e0, to=e1, relation_label="R2"
-    )
-
-    bm.relations.extend([r0, r1, r2])
-
-    await processor.process(message=bm, seqid=1)
-
-    index = get_indexing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    pb = await storage.get_indexing(index._calls[0][1])
-
-    assert len(pb.relations) == 3
-    assert pb.relations[0] == r0
-    assert pb.relations[1] == r1
-    assert pb.relations[2] == r2
-
-
-@pytest.mark.asyncio
-async def test_ingest_label_relation_extraction(
-    fake_node, local_files, storage, knowledgebox_ingest, processor
-):
-    rid = str(uuid.uuid4())
-    bm = BrokerMessage(
-        kbid=knowledgebox_ingest, uuid=rid, slug="slug-1", type=BrokerMessage.AUTOCOMMIT
-    )
-
-    labels = [
-        ("labelset-1", "label-1"),
-        ("labelset-1", "label-2"),
-        ("labelset-2", "label-1"),
-        ("labelset-2", "label-3"),
-    ]
-    bm.basic.usermetadata.classifications.extend(
-        [Classification(labelset=labelset, label=label) for labelset, label in labels]
-    )
-
-    await processor.process(message=bm, seqid=1)
-
-    index = get_indexing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    pb = await storage.get_indexing(index._calls[0][1])
-
-    for i, (labelset, label) in enumerate(labels):
-        assert pb.relations[i].relation == Relation.RelationType.ABOUT
-        assert pb.relations[i].source.value == rid
-        assert pb.relations[i].to.value == f"{labelset}/{label}"
-
-
-@pytest.mark.asyncio
-async def test_ingest_colab_relation_extraction(
-    fake_node, local_files, storage, knowledgebox_ingest, processor
-):
-    rid = str(uuid.uuid4())
-    bm = BrokerMessage(
-        kbid=knowledgebox_ingest, uuid=rid, slug="slug-1", type=BrokerMessage.AUTOCOMMIT
-    )
-
-    collaborators = ["Alice", "Bob", "Trudy"]
-    bm.origin.colaborators.extend(collaborators)
-
-    await processor.process(message=bm, seqid=1)
-
-    index = get_indexing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    pb = await storage.get_indexing(index._calls[0][1])
-
-    for i, collaborator in enumerate(collaborators):
-        assert pb.relations[i].relation == Relation.RelationType.COLAB
-        assert pb.relations[i].source.value == rid
-        assert pb.relations[i].to.value == collaborator
-
-
-@pytest.mark.asyncio
-async def test_ingest_field_metadata_relation_extraction(
-    fake_node, local_files, storage, knowledgebox_ingest, processor
-):
-    rid = str(uuid.uuid4())
-    bm = BrokerMessage(
-        kbid=knowledgebox_ingest,
-        uuid=rid,
-        slug="slug-1",
-        type=BrokerMessage.AUTOCOMMIT,
-        texts={
-            "title": FieldText(
-                body="Title with metadata",
-                format=FieldText.Format.PLAIN,
-            )
-        },
-    )
-
-    fcmw = FieldComputedMetadataWrapper(
-        field=FieldID(
-            field_type=FieldType.TEXT,
-            field="title",
-        )
-    )
-    fcmw.metadata.metadata.positions["subtype-1/value-1"].entity = "value-1"
-    fcmw.metadata.metadata.positions["subtype-1/value-2"].entity = "value-2"
-
-    fcmw.metadata.metadata.classifications.extend(
-        [
-            Classification(labelset="ls1", label="label1"),
-        ]
-    )
-
-    bm.field_metadata.append(fcmw)
-
-    await processor.process(message=bm, seqid=1)
-
-    index = get_indexing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    pb = await storage.get_indexing(index._calls[0][1])
-
-    generated_relations = [
-        # From ner metadata
-        Relation(
-            relation=Relation.RelationType.ENTITY,
-            source=RelationNode(value=rid, ntype=RelationNode.NodeType.RESOURCE),
-            to=RelationNode(
-                value="value-1", ntype=RelationNode.NodeType.ENTITY, subtype="subtype-1"
-            ),
-        ),
-        Relation(
-            relation=Relation.RelationType.ENTITY,
-            source=RelationNode(value=rid, ntype=RelationNode.NodeType.RESOURCE),
-            to=RelationNode(
-                value="value-2", ntype=RelationNode.NodeType.ENTITY, subtype="subtype-1"
-            ),
-        ),
-        # From classification metadata
-        Relation(
-            relation=Relation.RelationType.ABOUT,
-            source=RelationNode(value=rid, ntype=RelationNode.NodeType.RESOURCE),
-            to=RelationNode(
-                value="ls1/label1",
-                ntype=RelationNode.NodeType.LABEL,
-            ),
-        ),
-    ]
-    for generated_relation in generated_relations:
-        assert generated_relation in pb.relations
-
-
-@pytest.mark.asyncio
-async def test_ingest_field_relations_relation_extraction(
-    fake_node, local_files, storage, knowledgebox_ingest, processor
-):
-    rid = str(uuid.uuid4())
-    bm = BrokerMessage(
-        kbid=knowledgebox_ingest, uuid=rid, slug="slug-1", type=BrokerMessage.AUTOCOMMIT
-    )
-
-    relationnode = RelationNode(
-        value=rid, ntype=RelationNode.NodeType.RESOURCE, subtype="subtype-1"
-    )
-    test_relations = [
-        Relation(
-            relation=Relation.RelationType.CHILD,
-            source=relationnode,
-            to=RelationNode(
-                value="document",
-                ntype=RelationNode.NodeType.RESOURCE,
-            ),
-        ),
-        Relation(
-            relation=Relation.RelationType.ABOUT,
-            source=relationnode,
-            to=RelationNode(
-                value="label",
-                ntype=RelationNode.NodeType.LABEL,
-            ),
-        ),
-        Relation(
-            relation=Relation.RelationType.ENTITY,
-            source=relationnode,
-            to=RelationNode(
-                value="entity",
-                ntype=RelationNode.NodeType.ENTITY,
-            ),
-        ),
-        Relation(
-            relation=Relation.RelationType.COLAB,
-            source=relationnode,
-            to=RelationNode(
-                value="user",
-                ntype=RelationNode.NodeType.USER,
-            ),
-        ),
-        Relation(
-            relation=Relation.RelationType.OTHER,
-            source=relationnode,
-            to=RelationNode(
-                value="other",
-                ntype=RelationNode.NodeType.RESOURCE,
-            ),
-        ),
-    ]
-    bm.relations.extend(test_relations)
-
-    await processor.process(message=bm, seqid=1)
-
-    index = get_indexing()
-    storage = await get_storage(service_name=SERVICE_NAME)
-
-    pb = await storage.get_indexing(index._calls[0][1])
-
-    assert len(pb.relations) == len(test_relations)
-    for relation in test_relations:
-        assert relation in pb.relations

nucliadb/ingest/tests/unit/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/unit/consumer/test_auditing.py

@@ -1,140 +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_protos.audit_pb2 import AuditKBCounter, AuditRequest
-from nucliadb_protos.writer_pb2 import BrokerMessage, Notification, ShardObject
-
-from nucliadb.ingest.consumer import auditing
-from nucliadb_protos import nodereader_pb2
-
-pytestmark = pytest.mark.asyncio
-
-
-@pytest.fixture()
-def pubsub():
-    mock = AsyncMock()
-    mock.parse = lambda x: x
-    yield mock
-
-
-@pytest.fixture()
-def reader():
-    yield AsyncMock()
-
-
-@pytest.fixture()
-def shard_manager(reader):
-    nm = MagicMock()
-    node = MagicMock(reader=reader)
-    nm.get_shards_by_kbid = AsyncMock(return_value=[ShardObject()])
-    with patch(
-        "nucliadb.ingest.consumer.auditing.get_shard_manager", return_value=nm
-    ), patch(
-        "nucliadb.ingest.consumer.auditing.choose_node",
-        return_value=(node, "shard_id"),
-    ):
-        yield nm
-
-
-@pytest.fixture()
-def audit():
-    yield AsyncMock()
-
-
-@pytest.fixture()
-async def index_audit_handler(pubsub, audit, shard_manager):
-    iah = auditing.IndexAuditHandler(
-        driver=AsyncMock(transaction=MagicMock(return_value=AsyncMock())),
-        audit=audit,
-        pubsub=pubsub,
-        check_delay=0.05,
-    )
-    await iah.initialize()
-    yield iah
-    await iah.finalize()
-
-
-@pytest.fixture()
-async def writes_audit_handler(pubsub, audit, shard_manager):
-    rwah = auditing.ResourceWritesAuditHandler(
-        driver=AsyncMock(transaction=MagicMock(return_value=AsyncMock())),
-        storage=AsyncMock(),
-        audit=audit,
-        pubsub=pubsub,
-    )
-    await rwah.initialize()
-    yield rwah
-    await rwah.finalize()
-
-
-async def test_handle_message(
-    index_audit_handler: auditing.IndexAuditHandler, reader, audit
-):
-    reader.GetShard.return_value = nodereader_pb2.Shard(fields=5, paragraphs=6)
-
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.INDEXED,
-    )
-    await index_audit_handler.handle_message(notif.SerializeToString())
-
-    await asyncio.sleep(0.06)
-
-    audit.report.assert_called_with(
-        kbid="kbid",
-        audit_type=AuditRequest.AuditType.INDEXED,
-        kb_counter=AuditKBCounter(fields=5, paragraphs=6),
-    )
-
-
-async def test_handle_message_ignore_not_indexed(
-    index_audit_handler: auditing.IndexAuditHandler, audit
-):
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.COMMIT,
-    )
-    await index_audit_handler.handle_message(notif.SerializeToString())
-
-    await index_audit_handler.finalize()
-
-    audit.report.assert_not_called()
-
-
-async def test_resource_handle_message_processor_messages_are_not_audited(
-    writes_audit_handler: auditing.ResourceWritesAuditHandler, audit
-):
-    message = BrokerMessage()
-    message.source = BrokerMessage.MessageSource.PROCESSOR
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.COMMIT,
-        message=message,
-        write_type=Notification.WriteType.MODIFIED,
-    )
-    await writes_audit_handler.handle_message(notif.SerializeToString())
-
-    await writes_audit_handler.finalize()
-
-    audit.report.assert_not_called()

nucliadb/ingest/tests/unit/consumer/test_consumer.py

@@ -1,69 +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 pytest
-from nucliadb_protos.writer_pb2 import BrokerMessage, BrokerMessageBlobReference
-
-from nucliadb.ingest.consumer.consumer import IngestConsumer
-
-
-@pytest.fixture()
-def storage():
-    mock = MagicMock()
-    mock.get_stream_message = AsyncMock()
-    mock.del_stream_message = AsyncMock()
-    yield mock
-
-
-@pytest.fixture()
-def consumer(storage):
-    yield IngestConsumer(None, "partition", storage, None)
-
-
-@pytest.mark.asyncio
-async def test_get_broker_message(consumer: IngestConsumer, storage):
-    bm = BrokerMessage(kbid="kbid")
-    msg = Mock(data=bm.SerializeToString(), headers={})
-    assert bm == await consumer.get_broker_message(msg)
-    storage.get_stream_message.assert_not_called()
-
-
-@pytest.mark.asyncio
-async def test_get_broker_message_proxied(consumer: IngestConsumer, storage):
-    bm = BrokerMessage(kbid="kbid")
-    bmr = BrokerMessageBlobReference(kbid="kbid", storage_key="storage_key")
-    msg = Mock(data=bmr.SerializeToString(), headers={"X-MESSAGE-TYPE": "PROXY"})
-
-    storage.get_stream_message.return_value = bm.SerializeToString()
-
-    assert bm == await consumer.get_broker_message(msg)
-
-    storage.get_stream_message.assert_awaited_once_with("storage_key")
-
-
-@pytest.mark.asyncio
-async def test_clean_broker_message_proxied(consumer: IngestConsumer, storage):
-    bmr = BrokerMessageBlobReference(kbid="kbid", storage_key="storage_key")
-    msg = Mock(data=bmr.SerializeToString(), headers={"X-MESSAGE-TYPE": "PROXY"})
-
-    await consumer.clean_broker_message(msg)
-
-    storage.del_stream_message.assert_awaited_once_with("storage_key")

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

@@ -1,60 +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, patch
-
-import pytest
-
-from nucliadb.ingest.consumer.pull import PullWorker
-
-
-class TestPullWorker:
-    """
-    It's a complex class so this might get a little messy with mocks
-
-    It should be refactor at some point and these tests be rewritten/removed
-    """
-
-    @pytest.fixture()
-    def processor(self):
-        processor = AsyncMock()
-        with patch("nucliadb.ingest.consumer.pull.Processor", return_value=processor):
-            yield processor
-
-    @pytest.fixture()
-    def nats_conn(self):
-        conn = MagicMock()
-        conn.jetstream.return_value = AsyncMock()
-        conn.drain = AsyncMock()
-        conn.close = AsyncMock()
-        with patch("nucliadb.ingest.consumer.pull.nats.connect", return_value=conn):
-            yield conn
-
-    @pytest.fixture()
-    def worker(self, processor):
-        yield PullWorker(
-            driver=AsyncMock(),
-            partition="1",
-            storage=AsyncMock(),
-            pull_time_error_backoff=100,
-            zone="zone",
-            nuclia_processing_cluster_url="nuclia_processing_cluster_url",
-            nuclia_public_url="nuclia_public_url",
-            audit=None,
-            onprem=False,
-        )