diff --git a/swh/storage/in_memory.py b/swh/storage/in_memory.py --- a/swh/storage/in_memory.py +++ b/swh/storage/in_memory.py @@ -460,6 +460,7 @@ 'branches': copy.deepcopy(snapshot['branches']), '_sorted_branch_names': sorted(snapshot['branches']) } + self._objects[snapshot_id].append(('snapshot', snapshot_id)) self._origin_visits[visit]['snapshot'] = snapshot_id def snapshot_get(self, snapshot_id): @@ -780,6 +781,7 @@ origin['visits_dates'] = defaultdict(set) if key not in self._origins: self._origins[key] = origin + self._objects[key].append(('origin', key)) return key def origin_visit_add(self, origin, date=None, *, ts=None): @@ -914,7 +916,9 @@ ) stats = {key: 0 for key in keys} stats.update(collections.Counter( - obj_type for (obj_type, obj_id) in self._objects.values())) + obj_type + for (obj_type, obj_id) + in itertools.chain(*self._objects.values()))) return stats def refresh_stat_counters(self): diff --git a/swh/storage/tests/test_storage.py b/swh/storage/tests/test_storage.py --- a/swh/storage/tests/test_storage.py +++ b/swh/storage/tests/test_storage.py @@ -1414,12 +1414,45 @@ expected_keys = ['content', 'directory', 'directory_entry_dir', 'origin', 'person', 'revision'] + # Initially, all counters are 0 + self.storage.refresh_stat_counters() + counters = self.storage.stat_counters() + self.assertTrue(set(expected_keys) <= set(counters)) + for key in expected_keys: + self.assertEqual(counters[key], 0) + + # Add a content. Only the content counter should increase. + + self.storage.content_add([self.cont]) + self.storage.refresh_stat_counters() counters = self.storage.stat_counters() self.assertTrue(set(expected_keys) <= set(counters)) - self.assertIsInstance(counters[expected_keys[0]], int) + for key in expected_keys: + if key != 'content': + self.assertEqual(counters[key], 0) + self.assertEqual(counters['content'], 1) + + # Add other objects. Check their counter increased as well. + + origin_id = self.storage.origin_add_one(self.origin2) + origin_visit1 = self.storage.origin_visit_add( + origin_id, + date=self.date_visit2) + self.storage.snapshot_add(origin_id, origin_visit1['visit'], + self.snapshot) + self.storage.directory_add([self.dir]) + self.storage.revision_add([self.revision]) + + self.storage.refresh_stat_counters() + counters = self.storage.stat_counters() + self.assertEqual(counters['content'], 1) + self.assertEqual(counters['directory'], 1) + self.assertEqual(counters['snapshot'], 1) + self.assertEqual(counters['origin'], 1) + self.assertEqual(counters['revision'], 1) def test_content_find_with_present_content(self): # 1. with something to find