def testBestCheckpoints(self): """ Tests that the best checkpoints are tracked and ordered correctly. """ keep_checkpoints_num = 4 checkpoint_manager = CheckpointManager(keep_checkpoints_num, "i") checkpoints = [ Checkpoint(Checkpoint.MEMORY, i, self.mock_result(i)) for i in range(16) ] random.shuffle(checkpoints) for checkpoint in checkpoints: checkpoint_manager.on_checkpoint(checkpoint) best_checkpoints = checkpoint_manager.best_checkpoints() self.assertEqual(len(best_checkpoints), keep_checkpoints_num) for i in range(len(best_checkpoints)): self.assertEqual(best_checkpoints[i].value, i + 12)
def testOnCheckpointUnordered(self): """ Tests priorities that aren't inserted in ascending order. Also tests that the worst checkpoints are deleted when necessary. """ keep_checkpoints_num = 2 checkpoint_manager = CheckpointManager(keep_checkpoints_num, "i") checkpoints = [ Checkpoint(Checkpoint.DISK, {i}, self.mock_result(i)) for i in range(3, -1, -1) ] with patch("shutil.rmtree") as rmtree_mock, patch("os.path"): for j in range(0, len(checkpoints)): checkpoint_manager.on_checkpoint(checkpoints[j]) expected_deletes = 0 if j != 3 else 1 self.assertEqual(rmtree_mock.call_count, expected_deletes) self.assertEqual(checkpoint_manager.newest_checkpoint, checkpoints[j]) best_checkpoints = checkpoint_manager.best_checkpoints() self.assertEqual(len(best_checkpoints), keep_checkpoints_num) self.assertIn(checkpoints[0], best_checkpoints) self.assertIn(checkpoints[1], best_checkpoints)