diff --git a/swh/journal/tests/test_inmemory.py b/swh/journal/tests/test_inmemory.py index c414a9b..009cdb5 100644 --- a/swh/journal/tests/test_inmemory.py +++ b/swh/journal/tests/test_inmemory.py @@ -1,48 +1,41 @@ +# Copyright (C) 2019-2022 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + import pytest from swh.journal.writer import model_object_dict_sanitizer from swh.journal.writer.inmemory import InMemoryJournalWriter from swh.model.model import BaseModel from swh.model.tests.swh_model_data import TEST_OBJECTS def test_write_additions_with_test_objects(): writer = InMemoryJournalWriter[BaseModel]( value_sanitizer=model_object_dict_sanitizer ) expected = [] + priv_expected = [] for object_type, objects in TEST_OBJECTS.items(): writer.write_additions(object_type, objects) for object in objects: - expected.append((object_type, object)) + if object.anonymize(): + expected.append((object_type, object.anonymize())) + priv_expected.append((object_type, object)) + else: + expected.append((object_type, object)) - assert list(writer.privileged_objects) == [] + assert set(priv_expected) == set(writer.privileged_objects) assert set(expected) == set(writer.objects) -def test_write_additions_with_privileged_test_objects(): - writer = InMemoryJournalWriter[BaseModel]( - value_sanitizer=model_object_dict_sanitizer - ) - - expected = [] - - for object_type, objects in TEST_OBJECTS.items(): - writer.write_additions(object_type, objects, True) - - for object in objects: - expected.append((object_type, object)) - - assert list(writer.objects) == [] - assert set(expected) == set(writer.privileged_objects) - - def test_write_addition_errors_without_unique_key(): writer = InMemoryJournalWriter[BaseModel]( value_sanitizer=model_object_dict_sanitizer ) with pytest.raises(NotImplementedError): writer.write_addition("BaseModel", BaseModel()) diff --git a/swh/journal/writer/inmemory.py b/swh/journal/writer/inmemory.py index 69f63f8..775785a 100644 --- a/swh/journal/writer/inmemory.py +++ b/swh/journal/writer/inmemory.py @@ -1,45 +1,43 @@ -# Copyright (C) 2019 The Software Heritage developers +# Copyright (C) 2019-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information import logging from multiprocessing import Manager from typing import Any, Callable, Dict, Generic, List, Tuple, TypeVar from . import ValueProtocol logger = logging.getLogger(__name__) TValue = TypeVar("TValue", bound=ValueProtocol) class InMemoryJournalWriter(Generic[TValue]): objects: List[Tuple[str, TValue]] privileged_objects: List[Tuple[str, TValue]] def __init__( self, value_sanitizer: Callable[[str, Dict[str, Any]], Dict[str, Any]] ): # Share the list of objects across processes, for RemoteAPI tests. self.manager = Manager() self.objects = self.manager.list() self.privileged_objects = self.manager.list() - def write_addition( - self, object_type: str, object_: TValue, privileged: bool = False - ) -> None: + def write_addition(self, object_type: str, object_: TValue) -> None: object_.unique_key() # Check this does not error, to mimic the kafka writer - if privileged: + anon_object_ = object_.anonymize() + if anon_object_ is not None: self.privileged_objects.append((object_type, object_)) + self.objects.append((object_type, anon_object_)) else: self.objects.append((object_type, object_)) write_update = write_addition - def write_additions( - self, object_type: str, objects: List[TValue], privileged: bool = False - ) -> None: + def write_additions(self, object_type: str, objects: List[TValue]) -> None: for object_ in objects: - self.write_addition(object_type, object_, privileged) + self.write_addition(object_type, object_) diff --git a/swh/journal/writer/stream.py b/swh/journal/writer/stream.py index 202e13c..5d868b8 100644 --- a/swh/journal/writer/stream.py +++ b/swh/journal/writer/stream.py @@ -1,47 +1,43 @@ -# Copyright (C) 2021 The Software Heritage developers +# Copyright (C) 2021-2022 The Software Heritage developers # See the AUTHORS file at the top-level directory of this distribution # License: GNU General Public License version 3, or any later version # See top-level LICENSE file for more information import logging from typing import Any, BinaryIO, Callable, Dict, Generic, List, TypeVar from swh.journal.serializers import value_to_kafka from . import ValueProtocol logger = logging.getLogger(__name__) TValue = TypeVar("TValue", bound=ValueProtocol) class StreamJournalWriter(Generic[TValue]): """A simple JournalWriter which serializes objects in a stream Might be used to serialize a storage in a file to generate a test dataset. """ def __init__( self, output_stream: BinaryIO, value_sanitizer: Callable[[str, Dict[str, Any]], Dict[str, Any]], ): # Share the list of objects across processes, for RemoteAPI tests. self.output = output_stream self.value_sanitizer = value_sanitizer - def write_addition( - self, object_type: str, object_: TValue, privileged: bool = False - ) -> None: + def write_addition(self, object_type: str, object_: TValue) -> None: object_.unique_key() # Check this does not error, to mimic the kafka writer dict_ = self.value_sanitizer(object_type, object_.to_dict()) self.output.write(value_to_kafka((object_type, dict_))) write_update = write_addition - def write_additions( - self, object_type: str, objects: List[TValue], privileged: bool = False - ) -> None: + def write_additions(self, object_type: str, objects: List[TValue]) -> None: for object_ in objects: - self.write_addition(object_type, object_, privileged) + self.write_addition(object_type, object_)