def test_cleanup_serialization(self): self.manager = DependencyManager(self.work_dir, 5 * 1024 * 1024, 50) self.manager.finish_run( 'uuid1') # Has dependency, so will not be removed. self.manager.add_dependency('uuid1', '', 'uuid100') self.manager.add_dependency( 'uuid2', '', 'uuid100') # Downloading, so will not be removed. self.manager.finish_run('uuid3') # Used after uuid4, so will be left. self.manager.finish_run('uuid4') # Will be removed. self.manager.add_dependency('uuid4', '', 'uuid100') self.manager.remove_dependency('uuid4', '', 'uuid100') self.manager.add_dependency('uuid3', '', 'uuid100') self.manager.remove_dependency('uuid3', '', 'uuid100') for uuid in ['uuid1', 'uuid2', 'uuid3', 'uuid4']: with open(self.manager.get_run_path(uuid), 'wb') as f: f.write(' ' * 1024 * 1024) self.manager._cleanup_sleep_secs = 0 self.manager.start_cleanup_thread() time.sleep(0.1) self.manager.stop_cleanup_thread() self.check_state([('uuid1', ''), ('uuid3', '')]) self.assertIn(('uuid2', ''), self.manager._dependencies)
def test_cleanup(self): self.manager = DependencyManager(self.work_dir, 2 * 1024 * 1024, None) self.manager.finish_run( 'uuid1') # Has dependency, so will not be removed. self.manager.add_dependency('uuid1', '', 'uuid100') self.manager.add_dependency( 'uuid2', '', 'uuid100') # Downloading, so will not be removed. self.manager.finish_run('uuid3') # Used after uuid4, so will be left. self.manager.finish_run('uuid4') # Will be removed. self.manager.add_dependency('uuid4', '', 'uuid100') self.manager.remove_dependency('uuid4', '', 'uuid100') self.manager.add_dependency('uuid3', '', 'uuid100') self.manager.remove_dependency('uuid3', '', 'uuid100') for uuid in ['uuid1', 'uuid2', 'uuid3', 'uuid4']: with open(self.manager.get_run_path(uuid), 'wb') as f: f.write(' ' * 1024 * 1024) self.manager._cleanup_sleep_secs = 0 self.manager.start_cleanup_thread() time.sleep(0.1) self.manager.stop_cleanup_thread() self.check_state([('uuid1', ''), ('uuid3', '')]) self.assertIn(('uuid2', ''), self.manager._dependencies) self.assertItemsEqual([DependencyManager.STATE_FILENAME, 'bundles'], os.listdir(self.work_dir)) self.assertItemsEqual(['uuid1', 'uuid2', 'uuid3'], os.listdir(self.bundles_dir))
def test_load_state(self): self.manager.finish_run('uuid1') self.manager.finish_run('uuid2') with open(os.path.join(self.work_dir, 'random_file'), 'w'): pass self.assertIn('random_file', os.listdir(self.work_dir)) new_manager = DependencyManager(self.work_dir, 1 * 1024 * 1024) self.check_state([('uuid1', ''), ('uuid2', '')], new_manager) self.assertIn('state.json', os.listdir(self.work_dir)) self.assertNotIn('random_file', os.listdir(self.work_dir))
def setUp(self): self.work_dir = tempfile.mkdtemp() self.manager = DependencyManager(self.work_dir, None, None) self.bundles_dir = os.path.join(self.work_dir, 'bundles')
class DependencyManagerTest(unittest.TestCase): def setUp(self): self.work_dir = tempfile.mkdtemp() self.manager = DependencyManager(self.work_dir, None, None) self.bundles_dir = os.path.join(self.work_dir, 'bundles') def tearDown(self): remove_path(self.work_dir) def test_load_state(self): self.manager.finish_run('uuid1') self.manager.finish_run('uuid2') with open(os.path.join(self.bundles_dir, 'random_file'), 'w'): pass self.assertIn('random_file', os.listdir(self.bundles_dir)) new_manager = DependencyManager(self.work_dir, 1 * 1024 * 1024, None) self.check_state([('uuid1', ''), ('uuid2', '')], new_manager) self.assertIn(DependencyManager.STATE_FILENAME, os.listdir(self.work_dir)) self.assertNotIn('random_file', os.listdir(self.bundles_dir)) def test_downloading(self): self.manager.add_dependency('uuid1', '', 'uuid2') # Check cases that should not block. self.check_add_dependency_blocks('uuid1', 'a', 'uuid5', False, True) # Different path self.check_add_dependency_blocks('uuid6', '', 'uuid7', False, True) # Different UUID # This call will block until the failed download. Then, it will be asked # to retry the download. self.check_add_dependency_blocks('uuid1', '', 'uuid3', True, True) self.manager.finish_download('uuid1', '', False) self.check_state([]) time.sleep(0.01) # This call will block until both attempts to download are done. Then, # it will not be asked to retry the download. self.check_add_dependency_blocks('uuid1', '', 'uuid4', True, False) self.manager.finish_download('uuid1', '', True) self.check_state([('uuid1', '')]) def test_download_path_conflict(self): self.assertEqual( self.manager.add_dependency('uuid1', 'a/b_c', 'uuid2')[0], os.path.join(self.bundles_dir, 'uuid1_a_b_c')) self.assertEqual( self.manager.add_dependency('uuid1', 'a_b/c', 'uuid2')[0], os.path.join(self.bundles_dir, 'uuid1_a_b_c_')) def test_cleanup(self): self.manager = DependencyManager(self.work_dir, 2 * 1024 * 1024, None) self.manager.finish_run( 'uuid1') # Has dependency, so will not be removed. self.manager.add_dependency('uuid1', '', 'uuid100') self.manager.add_dependency( 'uuid2', '', 'uuid100') # Downloading, so will not be removed. self.manager.finish_run('uuid3') # Used after uuid4, so will be left. self.manager.finish_run('uuid4') # Will be removed. self.manager.add_dependency('uuid4', '', 'uuid100') self.manager.remove_dependency('uuid4', '', 'uuid100') self.manager.add_dependency('uuid3', '', 'uuid100') self.manager.remove_dependency('uuid3', '', 'uuid100') for uuid in ['uuid1', 'uuid2', 'uuid3', 'uuid4']: with open(self.manager.get_run_path(uuid), 'wb') as f: f.write(' ' * 1024 * 1024) self.manager._cleanup_sleep_secs = 0 self.manager.start_cleanup_thread() time.sleep(0.1) self.manager.stop_cleanup_thread() self.check_state([('uuid1', ''), ('uuid3', '')]) self.assertIn(('uuid2', ''), self.manager._dependencies) self.assertItemsEqual([DependencyManager.STATE_FILENAME, 'bundles'], os.listdir(self.work_dir)) self.assertItemsEqual(['uuid1', 'uuid2', 'uuid3'], os.listdir(self.bundles_dir)) def test_cleanup_serialization(self): self.manager = DependencyManager(self.work_dir, 5 * 1024 * 1024, 50) self.manager.finish_run( 'uuid1') # Has dependency, so will not be removed. self.manager.add_dependency('uuid1', '', 'uuid100') self.manager.add_dependency( 'uuid2', '', 'uuid100') # Downloading, so will not be removed. self.manager.finish_run('uuid3') # Used after uuid4, so will be left. self.manager.finish_run('uuid4') # Will be removed. self.manager.add_dependency('uuid4', '', 'uuid100') self.manager.remove_dependency('uuid4', '', 'uuid100') self.manager.add_dependency('uuid3', '', 'uuid100') self.manager.remove_dependency('uuid3', '', 'uuid100') for uuid in ['uuid1', 'uuid2', 'uuid3', 'uuid4']: with open(self.manager.get_run_path(uuid), 'wb') as f: f.write(' ' * 1024 * 1024) self.manager._cleanup_sleep_secs = 0 self.manager.start_cleanup_thread() time.sleep(0.1) self.manager.stop_cleanup_thread() self.check_state([('uuid1', ''), ('uuid3', '')]) self.assertIn(('uuid2', ''), self.manager._dependencies) def check_state(self, expected_targets, manager=None): if manager is None: manager = self.manager targets = [] expected_paths = [] with manager._lock: for target, dependency in manager._dependencies.iteritems(): if not dependency.downloading: targets.append(target) expected_paths.append(dependency.path) self.assertItemsEqual(expected_targets, targets) self.assertItemsEqual(expected_paths, manager._paths) state_file_targets = [] with open(manager._state_file, 'r') as f: for dep in json.loads(f.read()): state_file_targets.append(tuple(dep['target'])) self.assertItemsEqual(expected_targets, state_file_targets) def check_add_dependency_blocks(self, parent_uuid, parent_path, uuid, expected_blocks, expected_should_download): blocked = [True] def blocking_code(): _, should_download = self.manager.add_dependency( parent_uuid, parent_path, uuid) self.assertEqual(should_download, expected_should_download, msg='%s/%s from %s' % (parent_uuid, parent_path, uuid)) blocked[0] = False threading.Thread(target=blocking_code).start() time.sleep(0.01) self.assertEqual(blocked[0], expected_blocks)
def setUp(self): self.work_dir = tempfile.mkdtemp() self.manager = DependencyManager(self.work_dir, None)