diff --git a/swh/journal/tests/test_inmemory.py b/swh/journal/tests/test_inmemory.py --- a/swh/journal/tests/test_inmemory.py +++ b/swh/journal/tests/test_inmemory.py @@ -11,8 +11,10 @@ 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 = [] @@ -30,6 +32,20 @@ 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) diff --git a/swh/journal/writer/inmemory.py b/swh/journal/writer/inmemory.py --- a/swh/journal/writer/inmemory.py +++ b/swh/journal/writer/inmemory.py @@ -17,16 +17,21 @@ 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_))