def test_delete_process_checkpoints(self):
"""
"""
process_a = ProcessWithCheckpoint()
process_b = ProcessWithCheckpoint()
checkpoint_a1 = PersistedCheckpoint(process_a.pid, '1')
checkpoint_a2 = PersistedCheckpoint(process_a.pid, '2')
checkpoint_b1 = PersistedCheckpoint(process_b.pid, '1')
checkpoint_b2 = PersistedCheckpoint(process_b.pid, '2')
with tempfile.TemporaryDirectory() as directory:
persister = PicklePersister(directory)
persister.save_checkpoint(process_a, tag='1')
persister.save_checkpoint(process_a, tag='2')
persister.save_checkpoint(process_b, tag='1')
persister.save_checkpoint(process_b, tag='2')
checkpoints = [checkpoint_a1, checkpoint_a2]
retrieved_checkpoints = persister.get_process_checkpoints(
process_a.pid)
self.assertEquals(set(retrieved_checkpoints), set(checkpoints))
persister.delete_process_checkpoints(process_a.pid)
checkpoints = []
retrieved_checkpoints = persister.get_process_checkpoints(
process_a.pid)
self.assertEquals(set(retrieved_checkpoints), set(checkpoints))
def test_load_all_checkpoints(self):
self._empty_directory()
# Create some processes
for i in range(0, 3):
proc = ProcessWithCheckpoint.new(pid=i)
self.pickle_persistence.save(proc)
# Check that the number of checkpoints matches we expected
num_cps = len(self.pickle_persistence.load_all_checkpoints())
self.assertEqual(num_cps, 3)
def test_get_checkpoints_without_tags(self):
"""
"""
process_a = ProcessWithCheckpoint()
process_b = ProcessWithCheckpoint()
checkpoint_a = PersistedCheckpoint(process_a.pid, None)
checkpoint_b = PersistedCheckpoint(process_b.pid, None)
checkpoints = [checkpoint_a, checkpoint_b]
with tempfile.TemporaryDirectory() as directory:
persister = PicklePersister(directory)
persister.save_checkpoint(process_a)
persister.save_checkpoint(process_b)
retrieved_checkpoints = persister.get_checkpoints()
self.assertEquals(set(retrieved_checkpoints), set(checkpoints))
def test_save_load_roundtrip(self):
"""
Test the PicklePersister by taking a dummpy process, saving a checkpoint
and recreating it from the same checkpoint
"""
process = ProcessWithCheckpoint()
with tempfile.TemporaryDirectory() as directory:
persister = PicklePersister(directory)
persister.save_checkpoint(process)
bundle = persister.load_checkpoint(process.pid)
recreated = bundle.unbundle(loop=self.loop)
def test_on_finishing_process(self):
proc = ProcessWithCheckpoint.new()
pid = proc.pid
self.pickle_persistence.persist_process(proc)
running_path = self.pickle_persistence.get_running_path(proc.pid)
self.assertTrue(os.path.isfile(running_path))
proc.play()
self.assertFalse(os.path.isfile(running_path))
finished_path = \
os.path.join(self.store_dir,
self.pickle_persistence.finished_directory,
self.pickle_persistence.pickle_filename(pid))
self.assertTrue(os.path.isfile(finished_path))
def test_save(self):
proc = ProcessWithCheckpoint.new()
running_path = self.pickle_persistence.get_running_path(proc.pid)
self.pickle_persistence.save(proc)
self.assertTrue(os.path.isfile(running_path))
def test_on_create_process(self):
proc = ProcessWithCheckpoint.new()
self.pickle_persistence.persist_process(proc)
save_path = self.pickle_persistence.get_running_path(proc.pid)
self.assertTrue(os.path.isfile(save_path))