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

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

@@ -1,139 +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.writer_pb2 import Notification, ShardObject, Shards
-
-from nucliadb.common.cluster.settings import settings
-from nucliadb.ingest.consumer import shard_creator
-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 kbdm():
-    mock = MagicMock()
-    mock.get_model_metadata = AsyncMock(return_value="model")
-    with patch("nucliadb.common.cluster.manager.datamanagers.kb", return_value=mock):
-        yield mock
-
-
-@pytest.fixture()
-def shard_manager(reader):
-    sm = MagicMock()
-    node = MagicMock(reader=reader)
-    shards = Shards(shards=[ShardObject(read_only=False)], actual=0)
-    sm.get_current_active_shard = AsyncMock(return_value=shards.shards[0])
-    sm.maybe_create_new_shard = AsyncMock()
-    with patch(
-        "nucliadb.ingest.consumer.shard_creator.get_shard_manager", return_value=sm
-    ), patch(
-        "nucliadb.ingest.consumer.shard_creator.choose_node",
-        return_value=(node, "shard_id"),
-    ), patch(
-        "nucliadb.ingest.consumer.shard_creator.datamanagers.cluster.get_kb_shards",
-        AsyncMock(return_value=shards),
-    ), patch(
-        "nucliadb.ingest.consumer.shard_creator.locking.distributed_lock",
-        return_value=AsyncMock(),
-    ):
-        yield sm
-
-
-@pytest.fixture()
-async def shard_creator_handler(pubsub, shard_manager):
-    sc = shard_creator.ShardCreatorHandler(
-        driver=AsyncMock(transaction=MagicMock(return_value=AsyncMock())),
-        storage=AsyncMock(),
-        pubsub=pubsub,
-        check_delay=0.05,
-    )
-    await sc.initialize()
-    yield sc
-    await sc.finalize()
-
-
-async def test_handle_message_create_new_shard(
-    shard_creator_handler: shard_creator.ShardCreatorHandler,
-    reader,
-    kbdm,
-    shard_manager,
-):
-    reader.GetShard.return_value = nodereader_pb2.Shard(
-        paragraphs=settings.max_shard_paragraphs + 1
-    )
-
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.INDEXED,
-    )
-    await shard_creator_handler.handle_message(notif.SerializeToString())
-    await asyncio.sleep(0.06)
-    shard_manager.maybe_create_new_shard.assert_called_with(
-        "kbid", settings.max_shard_paragraphs + 1, 0, 0
-    )
-
-
-async def test_handle_message_do_not_create(
-    shard_creator_handler: shard_creator.ShardCreatorHandler, reader, shard_manager
-):
-    reader.GetShard.return_value = nodereader_pb2.Shard(
-        paragraphs=settings.max_shard_paragraphs - 1
-    )
-
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.INDEXED,
-    )
-    await shard_creator_handler.handle_message(notif.SerializeToString())
-
-    await shard_creator_handler.finalize()
-
-    shard_manager.create_shard_by_kbid.assert_not_called()
-
-
-async def test_handle_message_ignore_not_indexed(
-    shard_creator_handler: shard_creator.ShardCreatorHandler, shard_manager
-):
-    notif = Notification(
-        kbid="kbid",
-        action=Notification.Action.COMMIT,
-    )
-    await shard_creator_handler.handle_message(notif.SerializeToString())
-
-    await shard_creator_handler.finalize()
-
-    shard_manager.create_shard_by_kbid.assert_not_called()

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

@@ -1,67 +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 utils
-
-pytestmark = pytest.mark.asyncio
-
-
-async def test_delay_task_handler():
-    dth = utils.DelayedTaskHandler(0.05)
-    await dth.initialize()
-
-    counter = 0
-
-    async def handler():
-        await asyncio.sleep(0.1)
-        nonlocal counter
-        counter += 1
-
-    dth.schedule("key1", handler)
-    dth.schedule("key1", handler)
-    dth.schedule("key1", handler)
-    dth.schedule("key2", handler)
-    dth.schedule("key3", handler)
-    dth.schedule("key4", handler)
-
-    # all should be scheduled and duplicates ignored
-    assert len(dth.to_process) == 4
-
-    await asyncio.sleep(0.06)
-    # they should all be running now
-    assert len(dth.outstanding_tasks) == 4
-
-    # schedule a couple more
-    dth.schedule("key1", handler)  # duplicate key, should get rescheduled at end
-    dth.schedule("key5", handler)
-    dth.schedule("key6", handler)
-
-    await asyncio.sleep(0.1)
-    # original set should be finished now
-    assert counter == 4
-
-    # finish everything now
-    await dth.finalize()
-
-    assert counter == 7

nucliadb/ingest/tests/unit/orm/__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/ingest/tests/unit/orm/test_brain.py

@@ -1,247 +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_protos.noderesources_pb2 import Resource as PBResource
-from nucliadb_protos.resources_pb2 import (
-    Basic,
-    ExtractedText,
-    FieldComputedMetadataWrapper,
-    FieldID,
-    FieldType,
-    Metadata,
-    Paragraph,
-    Sentence,
-)
-
-from nucliadb.ingest.orm.brain import ParagraphPages, ResourceBrain
-from nucliadb_protos import resources_pb2
-
-
-def test_apply_field_metadata_marks_duplicated_paragraphs():
-    # Simulate a field with two paragraphs that contain the same text
-    br = ResourceBrain(rid=str(uuid4()))
-    field_key = "text1"
-    fcmw = FieldComputedMetadataWrapper()
-    fcmw.field.CopyFrom(FieldID(field_type=FieldType.TEXT, field=field_key))
-    paragraph = "Some paragraph here. "
-    text_1 = f"{paragraph}{paragraph}"
-    first_occurrence = [0, len(paragraph)]
-    second_occurrence = [len(paragraph), len(paragraph) * 2]
-
-    et = ExtractedText(text=text_1)
-    p1 = Paragraph(start=first_occurrence[0], end=first_occurrence[1])
-    p1.sentences.append(
-        Sentence(start=first_occurrence[0], end=first_occurrence[1], key="test")
-    )
-    p2 = Paragraph(start=second_occurrence[0], end=second_occurrence[1])
-    p2.sentences.append(
-        Sentence(start=second_occurrence[0], end=second_occurrence[1], key="test")
-    )
-    fcmw.metadata.metadata.paragraphs.append(p1)
-    fcmw.metadata.metadata.paragraphs.append(p2)
-
-    br.apply_field_metadata(
-        field_key,
-        fcmw.metadata,
-        replace_field=[],
-        replace_splits={},
-        page_positions={},
-        extracted_text=et,
-    )
-
-    assert len(br.brain.paragraphs[field_key].paragraphs) == 2
-    for key, paragraph in br.brain.paragraphs[field_key].paragraphs.items():
-        if f"{first_occurrence[0]}-{first_occurrence[1]}" in key:
-            # Only the first time that a paragraph is found should be set to false
-            assert paragraph.repeated_in_field is False
-        else:
-            assert paragraph.repeated_in_field is True
-
-
-def test_apply_field_metadata_marks_duplicated_paragraphs_on_split_metadata():
-    # # Test now the split text path
-    br = ResourceBrain(rid=str(uuid4()))
-    field_key = "text1"
-    split_key = "subfield"
-    fcmw = FieldComputedMetadataWrapper()
-    fcmw.field.CopyFrom(FieldID(field_type=FieldType.TEXT, field=field_key))
-    paragraph = "Some paragraph here. "
-    text_1 = f"{paragraph}{paragraph}"
-    first_occurrence = [0, len(paragraph)]
-    second_occurrence = [len(paragraph), len(paragraph) * 2]
-
-    et = ExtractedText()
-    et.split_text[split_key] = text_1
-    p1 = Paragraph(start=first_occurrence[0], end=first_occurrence[1])
-    p1.sentences.append(
-        Sentence(start=first_occurrence[0], end=first_occurrence[1], key="test")
-    )
-    p2 = Paragraph(start=second_occurrence[0], end=second_occurrence[1])
-    p2.sentences.append(
-        Sentence(start=second_occurrence[0], end=second_occurrence[1], key="test")
-    )
-    fcmw.metadata.split_metadata[split_key].paragraphs.append(p1)
-    fcmw.metadata.split_metadata[split_key].paragraphs.append(p2)
-
-    br.apply_field_metadata(
-        field_key,
-        fcmw.metadata,
-        replace_field=[],
-        replace_splits={},
-        page_positions={},
-        extracted_text=et,
-    )
-
-    assert len(br.brain.paragraphs[field_key].paragraphs) == 2
-    for key, paragraph in br.brain.paragraphs[field_key].paragraphs.items():
-        if f"{first_occurrence[0]}-{first_occurrence[1]}" in key:
-            # Only the first time that a paragraph is found should be set to false
-            assert paragraph.repeated_in_field is False
-        else:
-            assert paragraph.repeated_in_field is True
-
-
-def test_get_page_number():
-    page_numbers = ParagraphPages(
-        {
-            0: (0, 99),
-            1: (100, 199),
-            2: (200, 299),
-        }
-    )
-    assert page_numbers.get(10) == 0
-    assert page_numbers.get(100) == 1
-    assert page_numbers.get(500) == 2
-
-
-@pytest.mark.parametrize(
-    "new_status,previous_status,expected_brain_status",
-    [
-        # No previous_status
-        (Metadata.Status.PENDING, None, PBResource.PENDING),
-        (Metadata.Status.PROCESSED, None, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, None, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, None, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, None, PBResource.PROCESSED),
-        # previous_status = PENDING
-        (Metadata.Status.PENDING, Metadata.Status.PENDING, PBResource.PENDING),
-        (Metadata.Status.PROCESSED, Metadata.Status.PENDING, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, Metadata.Status.PENDING, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, Metadata.Status.PENDING, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, Metadata.Status.PENDING, PBResource.PROCESSED),
-        # previous_status = PROCESSED
-        (Metadata.Status.PROCESSED, Metadata.Status.PROCESSED, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, Metadata.Status.PROCESSED, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, Metadata.Status.PROCESSED, PBResource.PROCESSED),
-        (Metadata.Status.PENDING, Metadata.Status.PROCESSED, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, Metadata.Status.PROCESSED, PBResource.PROCESSED),
-        # previous_status = ERROR
-        (Metadata.Status.PENDING, Metadata.Status.ERROR, PBResource.PROCESSED),
-        (Metadata.Status.PROCESSED, Metadata.Status.ERROR, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, Metadata.Status.ERROR, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, Metadata.Status.ERROR, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, Metadata.Status.ERROR, PBResource.PROCESSED),
-        # previous_status = BLOCKED
-        (Metadata.Status.PENDING, Metadata.Status.BLOCKED, PBResource.PROCESSED),
-        (Metadata.Status.PROCESSED, Metadata.Status.BLOCKED, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, Metadata.Status.BLOCKED, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, Metadata.Status.BLOCKED, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, Metadata.Status.BLOCKED, PBResource.PROCESSED),
-        # previous_status = EXPIRED
-        (Metadata.Status.PENDING, Metadata.Status.EXPIRED, PBResource.PROCESSED),
-        (Metadata.Status.PROCESSED, Metadata.Status.EXPIRED, PBResource.PROCESSED),
-        (Metadata.Status.ERROR, Metadata.Status.EXPIRED, PBResource.PROCESSED),
-        (Metadata.Status.BLOCKED, Metadata.Status.EXPIRED, PBResource.PROCESSED),
-        (Metadata.Status.EXPIRED, Metadata.Status.EXPIRED, PBResource.PROCESSED),
-    ],
-)
-def test_set_processing_status(new_status, previous_status, expected_brain_status):
-    br = ResourceBrain(rid="foo")
-    basic = Basic()
-    basic.metadata.status = new_status
-    br.set_processing_status(basic, previous_status)
-    assert br.brain.status == expected_brain_status
-
-
-def test_apply_field_metadata_populates_page_number():
-    br = ResourceBrain(rid="foo")
-    field_key = "text1"
-
-    fcmw = FieldComputedMetadataWrapper()
-    fcmw.field.CopyFrom(FieldID(field_type=FieldType.TEXT, field=field_key))
-
-    p1 = Paragraph(
-        start=40, end=54, start_seconds=[0], end_seconds=[10], text="Some text here"
-    )
-    p1.sentences.append(Sentence(start=40, end=54, key="test"))
-    fcmw.metadata.metadata.paragraphs.append(p1)
-
-    # Add it to the split too
-    fcmw.metadata.split_metadata["subfield"].paragraphs.append(p1)
-
-    page_positions = {
-        0: (0, 20),
-        1: (21, 39),
-        2: (40, 100),
-    }
-    br.apply_field_metadata(
-        field_key,
-        fcmw.metadata,
-        replace_field=[],
-        replace_splits={},
-        page_positions=page_positions,
-        extracted_text=None,
-    )
-
-    assert len(br.brain.paragraphs[field_key].paragraphs) == 2
-    for paragraph in br.brain.paragraphs[field_key].paragraphs.values():
-        assert paragraph.metadata.position.page_number == 2
-        assert paragraph.metadata.position.start == 40
-        assert paragraph.metadata.position.end == 54
-        assert paragraph.metadata.position.start_seconds == [0]
-        assert paragraph.metadata.position.end_seconds == [10]
-
-
-def test_set_resource_metadata_promotes_origin_dates():
-    resource_brain = ResourceBrain("rid")
-    basic = Basic()
-    basic.created.seconds = 1
-    basic.modified.seconds = 2
-    origin = resources_pb2.Origin()
-    origin.created.seconds = 3
-    origin.modified.seconds = 4
-
-    resource_brain.set_resource_metadata(basic, origin)
-
-    assert resource_brain.brain.metadata.created.seconds == 3
-    assert resource_brain.brain.metadata.modified.seconds == 4
-
-
-def test_set_resource_metadata_handles_timestamp_not_present():
-    resource_brain = ResourceBrain("rid")
-    basic = Basic()
-    resource_brain.set_resource_metadata(basic, None)
-    created = resource_brain.brain.metadata.created.seconds
-    modified = resource_brain.brain.metadata.modified.seconds
-    assert created > 0
-    assert modified > 0
-    assert modified >= created

nucliadb/ingest/tests/unit/orm/test_processor.py

@@ -1,131 +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, patch
-
-import pytest
-
-from nucliadb.common.cluster.settings import settings as cluster_settings
-from nucliadb.ingest.orm.exceptions import ResourceNotIndexable
-from nucliadb.ingest.orm.processor import Processor, validate_indexable_resource
-from nucliadb_protos import noderesources_pb2
-
-
-@pytest.fixture()
-def txn():
-    yield AsyncMock()
-
-
-@pytest.fixture()
-def driver(txn):
-    mock = MagicMock()
-    mock.transaction.return_value.__aenter__.return_value = txn
-    yield mock
-
-
-@pytest.fixture()
-def sm():
-    mock = AsyncMock()
-    mock.add_resource = AsyncMock()
-    with patch("nucliadb.ingest.orm.processor.get_shard_manager", return_value=mock):
-        yield mock
-
-
-@pytest.fixture()
-def processor(driver, sm):
-    yield Processor(driver, None)
-
-
-@pytest.fixture()
-def resource():
-    mock = MagicMock()
-    mock.set_basic = AsyncMock()
-    yield mock
-
-
-@pytest.fixture()
-def kb():
-    mock = MagicMock(kbid="kbid")
-    mock.get_resource_shard_id = AsyncMock()
-    mock.get_resource_shard = AsyncMock()
-    yield mock
-
-
-async def test_commit_slug(processor: Processor, txn, resource):
-    another_txn = Mock()
-    resource.txn = another_txn
-    resource.set_slug = AsyncMock()
-
-    await processor.commit_slug(resource)
-
-    resource.set_slug.assert_awaited_once()
-    txn.commit.assert_awaited_once()
-    assert resource.txn is another_txn
-
-
-async def test_mark_resource_error(processor: Processor, txn, resource, kb, sm):
-    await processor._mark_resource_error(kb, resource, partition="partition", seqid=1)
-    txn.commit.assert_called_once()
-    resource.set_basic.assert_awaited_once()
-    sm.add_resource.assert_awaited_once_with(
-        kb.get_resource_shard.return_value,
-        resource.indexer.brain,
-        1,
-        partition="partition",
-        kb="kbid",
-    )
-
-
-async def test_mark_resource_error_handle_error(
-    processor: Processor, kb, resource, txn
-):
-    resource.set_basic.side_effect = Exception("test")
-    await processor._mark_resource_error(kb, resource, partition="partition", seqid=1)
-    txn.commit.assert_not_called()
-
-
-async def test_mark_resource_error_skip_no_shard(
-    processor: Processor, resource, driver, kb, txn
-):
-    kb.get_resource_shard.return_value = None
-    await processor._mark_resource_error(kb, resource, partition="partition", seqid=1)
-    txn.commit.assert_not_called()
-
-
-async def test_mark_resource_error_skip_no_resource(
-    processor: Processor, kb, driver, txn
-):
-    await processor._mark_resource_error(kb, None, partition="partition", seqid=1)
-    txn.commit.assert_not_called()
-
-
-def test_validate_indexable_resource():
-    resource = noderesources_pb2.Resource()
-    resource.paragraphs["test"].paragraphs["test"].sentences["test"].vector.append(1.0)
-    validate_indexable_resource(resource)
-
-
-def test_validate_indexable_resource_throws_error_for_max():
-    resource = noderesources_pb2.Resource()
-    for i in range(cluster_settings.max_resource_paragraphs + 1):
-        resource.paragraphs["test"].paragraphs[f"test{i}"].sentences[
-            "test"
-        ].vector.append(1.0)
-    with pytest.raises(ResourceNotIndexable):
-        validate_indexable_resource(resource)