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

nucliadb/tests/unit/test_purge.py

@@ -1,136 +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 import purge
-from nucliadb.common.cluster.exceptions import NodeError, ShardNotFound
-
-pytestmark = pytest.mark.asyncio
-
-
-class DataIterator:
-    def __init__(self, data):
-        self.data = data
-
-    def __call__(self, *args, **kwargs):
-        return self
-
-    async def __aiter__(self):
-        for item in self.data:
-            yield item
-
-
-@pytest.fixture
-def keys():
-    yield []
-
-
-@pytest.fixture
-def txn(keys):
-    mock = AsyncMock()
-    mock.keys = DataIterator(keys)
-    yield mock
-
-
-@pytest.fixture
-def driver(txn):
-    mock = AsyncMock()
-    cm = AsyncMock()
-    cm.__aenter__.return_value = txn
-    mock.transaction = MagicMock(return_value=cm)
-    yield mock
-
-
-@pytest.fixture
-def storage():
-    mock = AsyncMock()
-    mock.delete_kb.return_value = True, False
-    yield mock
-
-
-@pytest.fixture(autouse=True)
-def kb():
-    mock = AsyncMock()
-    with patch("nucliadb.purge.KnowledgeBox", mock):
-        yield mock
-
-
-async def test_purge(kb, keys, driver):
-    keys.append("/pathto/kbid")
-
-    await purge.purge_kb(driver)
-
-    kb.purge.assert_called_once_with(driver, "kbid")
-    driver.begin.return_value.commit.assert_called_once()
-
-
-async def test_purge_handle_errors(kb, keys, driver):
-    keys.append("/failed")
-    keys.append("/pathto/failed")
-    keys.append("/pathto/failed")
-    keys.append("/pathto/failed")
-    keys.append("/pathto/failed")
-
-    kb.purge.side_effect = [ShardNotFound(), NodeError(), Exception(), None]
-    driver.begin.return_value.delete.side_effect = Exception()
-
-    await purge.purge_kb(driver)
-
-    driver.begin.return_value.commit.assert_not_called()
-    driver.begin.return_value.abort.assert_called_once()
-
-
-async def test_purge_kb_storage(
-    keys,
-    driver,
-    storage,
-):
-    keys.append("/pathto/kbid")
-
-    await purge.purge_kb_storage(driver, storage)
-
-    driver.begin.return_value.commit.assert_called_once()
-
-
-async def test_purge_kb_storage_handle_errors(keys, driver, storage):
-    keys.append("/failed")
-    keys.append("/pathto/failed")
-
-    driver.begin.return_value.delete.side_effect = Exception()
-
-    await purge.purge_kb_storage(driver, storage)
-
-    driver.begin.return_value.commit.assert_not_called()
-
-
-async def test_main(driver, storage):
-    with (
-        patch("nucliadb.purge.purge_kb", AsyncMock()) as purge_kb,
-        patch("nucliadb.purge.purge_kb_storage", AsyncMock()) as purge_kb_storage,
-        patch("nucliadb.purge.get_storage", return_value=storage),
-        patch("nucliadb.purge.setup_driver", return_value=driver),
-        patch("nucliadb.purge.setup_cluster", return_value=driver),
-    ):
-        await purge.main()
-
-        purge_kb.assert_called_once_with(driver)
-        purge_kb_storage.assert_called_once_with(driver, storage)

nucliadb/tests/utils/__init__.py

@@ -1,74 +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
-from datetime import datetime
-
-from nucliadb_protos.writer_pb2 import BrokerMessage, OpStatusWriter
-from nucliadb_protos.writer_pb2_grpc import WriterStub
-
-from nucliadb_protos import resources_pb2 as rpb
-
-
-def broker_resource(
-    kbid: str, rid=None, slug=None, title=None, summary=None
-) -> BrokerMessage:
-    """
-    Returns a broker resource with barebones metadata.
-    """
-    rid = rid or str(uuid.uuid4())
-    slug = slug or f"{rid}slug1"
-    bm: BrokerMessage = BrokerMessage(
-        kbid=kbid,
-        uuid=rid,
-        slug=slug,
-        type=BrokerMessage.AUTOCOMMIT,
-    )
-    title = title or "Title Resource"
-    summary = summary or "Summary of document"
-    bm.basic.icon = "text/plain"
-    bm.basic.title = title
-    bm.basic.summary = summary
-    bm.basic.thumbnail = "doc"
-    bm.basic.layout = "default"
-    bm.basic.metadata.useful = True
-    bm.basic.metadata.language = "es"
-    bm.basic.created.FromDatetime(datetime.now())
-    bm.basic.modified.FromDatetime(datetime.now())
-    bm.origin.source = rpb.Origin.Source.WEB
-
-    etw = rpb.ExtractedTextWrapper()
-    etw.body.text = title
-    etw.field.field = "title"
-    etw.field.field_type = rpb.FieldType.GENERIC
-    bm.extracted_text.append(etw)
-
-    etw = rpb.ExtractedTextWrapper()
-    etw.body.text = summary
-    etw.field.field = "summary"
-    etw.field.field_type = rpb.FieldType.GENERIC
-    bm.extracted_text.append(etw)
-
-    bm.source = BrokerMessage.MessageSource.WRITER
-    return bm
-
-
-async def inject_message(writer: WriterStub, message: BrokerMessage):
-    resp = await writer.ProcessMessage([message])  # type: ignore
-    assert resp.status == OpStatusWriter.Status.OK

nucliadb/tests/utils/aiohttp_session.py

@@ -1,44 +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
-
-
-def get_mocked_session(
-    http_method: str, status: int, text=None, json=None, read=None, context_manager=True
-):
-    response = Mock(status=status)
-    if text is not None:
-        response.text = AsyncMock(return_value=text)
-    if json is not None:
-        response.json = AsyncMock(return_value=json)
-    if read is not None:
-        if isinstance(read, str):
-            read = read.encode()
-        response.read = AsyncMock(return_value=read)
-    if context_manager:
-        # For when async with self.session.post() as response: is called
-        session = Mock()
-        http_method_mock = AsyncMock(__aenter__=AsyncMock(return_value=response))
-        getattr(session, http_method.lower()).return_value = http_method_mock
-    else:
-        # For when await self.session.post() is called
-        session = AsyncMock()
-        getattr(session, http_method.lower()).return_value = response
-    return session

nucliadb/tests/utils/broker_messages/__init__.py

@@ -1,171 +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 typing import Optional
-from uuid import uuid4
-
-from nucliadb_protos import resources_pb2 as rpb
-from nucliadb_protos import writer_pb2 as wpb
-
-from .fields import FieldBuilder
-from .helpers import labels_to_classifications
-
-
-class BrokerMessageBuilder:
-    """Helper to deal with broker message creation. It allows customized
-    creation of broker messages with sensible defaults and default title and
-    summary.
-
-    """
-
-    def __init__(
-        self,
-        *,
-        kbid: str,
-        rid: Optional[str] = None,
-        slug: Optional[str] = None,
-    ):
-        self.bm = wpb.BrokerMessage()
-        self.fields: dict[tuple[str, rpb.FieldType.ValueType], FieldBuilder] = {}
-
-        self.bm.kbid = kbid
-        self.bm.type = wpb.BrokerMessage.AUTOCOMMIT
-
-        # if first BM comes from PROCESSOR, it'll be ignored as it's out of order
-        self.bm.source = wpb.BrokerMessage.MessageSource.WRITER
-
-        if rid is None:
-            rid = str(uuid4())
-        self.bm.uuid = rid
-
-        if slug is None:
-            slug = f"{rid}-slug"
-        self.bm.slug = slug
-
-        self._default_basic()
-        self._default_origin()
-
-    def build(self) -> wpb.BrokerMessage:
-        self._apply_fields()
-        return self.bm
-
-    def add_field_builder(self, field: FieldBuilder):
-        self.fields[(field.id.field, field.id.field_type)] = field
-
-    def field_builder(
-        self, field_id: str, field_type: rpb.FieldType.ValueType
-    ) -> FieldBuilder:
-        return self.fields[(field_id, field_type)]
-
-    def with_title(self, title: str):
-        title_builder = FieldBuilder("title", rpb.FieldType.GENERIC)
-        title_builder.with_extracted_text(title)
-        # we do this to writer BMs in write resource API endpoint
-        title_builder.with_extracted_paragraph_metadata(
-            rpb.Paragraph(
-                start=0,
-                end=len(title),
-                kind=rpb.Paragraph.TypeParagraph.TITLE,
-            )
-        )
-        self.bm.basic.title = title
-        self.add_field_builder(title_builder)
-
-    def with_summary(self, summary: str):
-        summary_builder = FieldBuilder("summary", rpb.FieldType.GENERIC)
-        summary_builder.with_extracted_text(summary)
-        # we do this to writer BMs in write resource API endpoint
-        summary_builder.with_extracted_paragraph_metadata(
-            rpb.Paragraph(
-                start=0,
-                end=len(summary),
-                kind=rpb.Paragraph.TypeParagraph.DESCRIPTION,
-            )
-        )
-        self.bm.basic.summary = summary
-        self.add_field_builder(summary_builder)
-
-    def with_resource_labels(self, labelset: str, labels: list[str]):
-        classifications = labels_to_classifications(labelset, labels)
-        self.bm.basic.usermetadata.classifications.extend(classifications)
-
-    def _default_basic(self):
-        self.bm.basic.icon = "text/plain"
-        self.bm.basic.thumbnail = "doc"
-        self.bm.basic.layout = "default"
-        self.bm.basic.metadata.useful = True
-        self.bm.basic.metadata.language = "en"
-        self.bm.basic.metadata.status = rpb.Metadata.Status.PROCESSED
-        self.bm.basic.metadata.metadata["key"] = "value"
-        self.bm.basic.created.FromDatetime(datetime.now())
-        self.bm.basic.modified.FromDatetime(datetime.now())
-
-        self.with_title("Default test resource title")
-        self.with_summary("Default test resource summary")
-
-    def _default_origin(self):
-        self.bm.origin.source = rpb.Origin.Source.API
-        self.bm.origin.source_id = "My Source"
-        self.bm.origin.created.FromDatetime(datetime.now())
-        self.bm.origin.modified.FromDatetime(datetime.now())
-
-    def _apply_fields(self):
-        def replace_if_exists(mut_iterable, field_id: rpb.FieldID, item):
-            for obj in mut_iterable:
-                if obj.field == field_id:
-                    obj.Clear()
-                    obj.CopyFrom(item)
-                    break
-            else:
-                mut_iterable.append(item)
-
-        for field_builder in self.fields.values():
-            field = field_builder.build()
-
-            if field.id.field_type == rpb.FieldType.GENERIC:
-                pass
-            elif field.id.field_type == rpb.FieldType.FILE:
-                file_field = self.bm.files[field.id.field]
-                file_field.added.FromDatetime(datetime.now())
-                file_field.file.source = rpb.CloudFile.Source.EXTERNAL
-            else:
-                raise Exception("Unsupported field type")
-
-            if field.user.metadata is not None:
-                replace_if_exists(
-                    self.bm.basic.fieldmetadata, field.id, field.user.metadata
-                )
-            if field.extracted.metadata is not None:
-                replace_if_exists(
-                    self.bm.field_metadata, field.id, field.extracted.metadata
-                )
-            if field.extracted.text is not None:
-                replace_if_exists(
-                    self.bm.extracted_text, field.id, field.extracted.text
-                )
-            if field.extracted.vectors is not None:
-                replace_if_exists(
-                    self.bm.field_vectors, field.id, field.extracted.vectors
-                )
-            if field.extracted.question_answers is not None:
-                replace_if_exists(
-                    self.bm.question_answers, field.id, field.extracted.question_answers
-                )

nucliadb/tests/utils/broker_messages/fields.py

@@ -1,197 +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 dataclasses
-from datetime import datetime
-from typing import Optional
-
-from nucliadb_protos import resources_pb2 as rpb
-from nucliadb_protos import utils_pb2
-
-from .helpers import labels_to_classifications
-
-
-@dataclasses.dataclass
-class FieldUser:
-    metadata: Optional[rpb.UserFieldMetadata] = None
-
-
-@dataclasses.dataclass
-class FieldExtracted:
-    metadata: Optional[rpb.FieldComputedMetadataWrapper] = None
-    text: Optional[rpb.ExtractedTextWrapper] = None
-    vectors: Optional[rpb.ExtractedVectorsWrapper] = None
-    question_answers: Optional[rpb.FieldQuestionAnswerWrapper] = None
-
-
-@dataclasses.dataclass
-class Field:
-    id: rpb.FieldID
-    user: FieldUser = dataclasses.field(default_factory=FieldUser)
-    extracted: FieldExtracted = dataclasses.field(default_factory=FieldExtracted)
-
-
-class FieldBuilder:
-    def __init__(self, field: str, field_type: rpb.FieldType.ValueType):
-        self._field_id = rpb.FieldID(field=field, field_type=field_type)
-        self.__extracted_metadata: Optional[rpb.FieldComputedMetadataWrapper] = None
-        self.__extracted_text: Optional[rpb.ExtractedTextWrapper] = None
-        self.__extracted_vectors: Optional[rpb.ExtractedVectorsWrapper] = None
-        self.__user_metadata: Optional[rpb.UserFieldMetadata] = None
-        self.__question_answers: Optional[rpb.FieldQuestionAnswerWrapper] = None
-
-    @property
-    def id(self) -> rpb.FieldID:
-        return self._field_id
-
-    # properties to generate a default value per pb
-
-    @property
-    def _extracted_metadata(self) -> rpb.FieldComputedMetadataWrapper:
-        if self.__extracted_metadata is None:
-            now = datetime.now()
-            self.__extracted_metadata = rpb.FieldComputedMetadataWrapper(
-                field=self._field_id,
-            )
-            self.__extracted_metadata.metadata.metadata.last_index.FromDatetime(now)
-            self.__extracted_metadata.metadata.metadata.last_understanding.FromDatetime(
-                now
-            )
-            self.__extracted_metadata.metadata.metadata.last_extract.FromDatetime(now)
-        return self.__extracted_metadata
-
-    @property
-    def _extracted_text(self) -> rpb.ExtractedTextWrapper:
-        if self.__extracted_text is None:
-            self.__extracted_text = rpb.ExtractedTextWrapper(field=self._field_id)
-        return self.__extracted_text
-
-    @property
-    def _extracted_vectors(self) -> rpb.ExtractedVectorsWrapper:
-        if self.__extracted_vectors is None:
-            self.__extracted_vectors = rpb.ExtractedVectorsWrapper(field=self._field_id)
-        return self.__extracted_vectors
-
-    @property
-    def _question_answers(self) -> rpb.FieldQuestionAnswerWrapper:
-        if self.__question_answers is None:
-            self.__question_answers = rpb.FieldQuestionAnswerWrapper(
-                field=self._field_id
-            )
-        return self.__question_answers
-
-    @property
-    def _user_metadata(self) -> rpb.UserFieldMetadata:
-        if self.__user_metadata is None:
-            self.__user_metadata = rpb.UserFieldMetadata(field=self._field_id)
-        return self.__user_metadata
-
-    def build(self) -> Field:
-        field = Field(id=self._field_id)
-
-        if self.__extracted_metadata is not None:
-            field.extracted.metadata = rpb.FieldComputedMetadataWrapper()
-            field.extracted.metadata.CopyFrom(self.__extracted_metadata)
-
-        if self.__extracted_text is not None:
-            field.extracted.text = rpb.ExtractedTextWrapper()
-            field.extracted.text.CopyFrom(self.__extracted_text)
-
-        if self.__extracted_vectors is not None:
-            field.extracted.vectors = rpb.ExtractedVectorsWrapper()
-            field.extracted.vectors.CopyFrom(self.__extracted_vectors)
-
-        if self.__question_answers is not None:
-            field.extracted.question_answers = rpb.FieldQuestionAnswerWrapper()
-            field.extracted.question_answers.CopyFrom(self.__question_answers)
-
-        if self.__user_metadata is not None:
-            field.user.metadata = rpb.UserFieldMetadata()
-            field.user.metadata.CopyFrom(self.__user_metadata)
-
-        return field
-
-    def with_extracted_labels(self, labelset: str, labels: list[str]):
-        classifications = labels_to_classifications(labelset, labels)
-        self._extracted_metadata.metadata.metadata.classifications.extend(
-            classifications
-        )
-
-    def with_extracted_text(self, text: str):
-        self._extracted_text.body.text = text
-
-    def with_extracted_vectors(self, vectors: list[utils_pb2.Vector]):
-        self._extracted_vectors.vectors.vectors.vectors.extend(vectors)
-
-    def with_extracted_paragraph_metadata(self, paragraph: rpb.Paragraph):
-        self._extracted_metadata.metadata.metadata.paragraphs.append(paragraph)
-
-    def with_user_entity(self, klass: str, name: str, *, start: int, end: int):
-        entity = rpb.TokenSplit(
-            klass=klass,
-            token=name,
-            start=start,
-            end=end,
-        )
-        self._user_metadata.token.append(entity)
-
-    def with_extracted_entity(
-        self, klass: str, name: str, *, positions: list[rpb.Position]
-    ):
-        entity = self._extracted_metadata.metadata.metadata.positions[f"{klass}/{name}"]
-        entity.entity = name
-        entity.position.extend(positions)
-
-    def with_user_paragraph_labels(self, key: str, labelset: str, labels: list[str]):
-        classifications = labels_to_classifications(labelset, labels)
-        pa = rpb.ParagraphAnnotation()
-        pa.key = key
-        pa.classifications.extend(classifications)
-        self._user_metadata.paragraphs.append(pa)
-
-    def add_question_answer(
-        self,
-        question: str,
-        answer: str,
-        question_lang: str = "en",
-        question_paragraph_ids: list[str] = [],
-        answer_lang: str = "en",
-        answer_paragraph_ids: list[str] = [],
-    ):
-        question_pb = rpb.Question(
-            text=question,
-            language=question_lang,
-            ids_paragraphs=question_paragraph_ids,
-        )
-        answer_pb = rpb.Answers(
-            text=answer,
-            language=answer_lang,
-            ids_paragraphs=answer_paragraph_ids,
-        )
-
-        # check if is another answer for an already added question
-        for question_answer in self._question_answers.question_answers.question_answer:
-            if question_answer.question == question_pb:
-                question_answer.answers.append(answer_pb)
-                return
-
-        question_answer = rpb.QuestionAnswer(question=question_pb)
-        question_answer.answers.append(answer_pb)
-        self._question_answers.question_answers.question_answer.append(question_answer)

nucliadb/tests/utils/broker_messages/helpers.py

@@ -1,33 +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_protos.resources_pb2 import Classification
-
-
-def labels_to_classifications(labelset: str, labels: list[str]) -> list[Classification]:
-    classifications = [
-        Classification(
-            labelset=labelset,
-            label=label,
-            cancelled_by_user=False,
-        )
-        for label in labels
-    ]
-    return classifications