diff --git a/swh/journal/tests/test_inmemory.py b/swh/journal/tests/test_inmemory.py index 8abd2d4..c1fa5f8 100644 --- a/swh/journal/tests/test_inmemory.py +++ b/swh/journal/tests/test_inmemory.py @@ -1,37 +1,53 @@ # 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(value_sanitizer=model_object_dict_sanitizer) +def test_write_additions_anonymized(): + writer = InMemoryJournalWriter( + value_sanitizer=model_object_dict_sanitizer, anonymize=True + ) expected = [] priv_expected = [] for object_type, objects in TEST_OBJECTS.items(): writer.write_additions(object_type, objects) for object in objects: if object.anonymize(): expected.append((object_type, object.anonymize())) priv_expected.append((object_type, object)) else: expected.append((object_type, object)) assert set(priv_expected) == set(writer.privileged_objects) assert set(expected) == set(writer.objects) +def test_write_additions(): + writer = InMemoryJournalWriter(value_sanitizer=model_object_dict_sanitizer) + expected = set() + + for object_type, objects in TEST_OBJECTS.items(): + writer.write_additions(object_type, objects) + + for object in objects: + expected.add((object_type, object)) + + assert not set(writer.privileged_objects) + assert expected == set(writer.objects) + + def test_write_addition_errors_without_unique_key(): writer = InMemoryJournalWriter(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 a2039c9..909ff46 100644 --- a/swh/journal/writer/inmemory.py +++ b/swh/journal/writer/inmemory.py @@ -1,43 +1,48 @@ # 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, Iterable, List, Tuple from .interface import ValueProtocol logger = logging.getLogger(__name__) class InMemoryJournalWriter: objects: List[Tuple[str, ValueProtocol]] privileged_objects: List[Tuple[str, ValueProtocol]] def __init__( - self, value_sanitizer: Callable[[str, Dict[str, Any]], Dict[str, Any]] + self, + value_sanitizer: Callable[[str, Dict[str, Any]], Dict[str, Any]], + anonymize: bool = False, ): # Share the list of objects across processes, for RemoteAPI tests. self.manager = Manager() self.objects = self.manager.list() self.privileged_objects = self.manager.list() + self.anonymize = anonymize def write_addition(self, object_type: str, object_: ValueProtocol) -> None: object_.unique_key() # Check this does not error, to mimic the kafka writer - anon_object_ = object_.anonymize() + anon_object_ = None + if self.anonymize: + 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_)) def write_additions( self, object_type: str, objects: Iterable[ValueProtocol] ) -> None: for object_ in objects: self.write_addition(object_type, object_) def flush(self) -> None: pass