def test_get_all_rank_n(self):
store = mock.MagicMock()
world_size = 3
store_util.get_all(store, 1, "test/store", world_size)
# omit empty kwargs, get only key
actual_set_call_args = [
call_args[0][0] for call_args in store.set.call_args_list
]
self.assertListEqual(["test/store1.FIN"], actual_set_call_args)
actual_get_call_args = [
call_args[0] for call_args in store.get.call_args_list
]
expected_get_call_args = [
("test/store0", ),
("test/store1", ),
("test/store2", ),
]
self.assertListEqual(expected_get_call_args, actual_get_call_args)
def test_get_data(self):
store = TestStore()
data = store_util.get_all(store, "test/store", 10)
for idx in range(0, 10):
self.assertEqual(f"retrieved:test/store{idx}", data[idx])