def testParseSpillDirs(self): self.assertEqual([], parse_spill_dirs(None)) self.assertEqual(['/tmp/a', '/tmp/b'], parse_spill_dirs(['/tmp/a', '/tmp/b'])) self.assertEqual([], parse_spill_dirs(os.path.pathsep)) temp_dir = tempfile.mkdtemp(prefix='test_mars_spill_') try: dirs = [ os.path.join(temp_dir, 'select_dir'), os.path.join(temp_dir, 'dir1'), os.path.join(temp_dir, 'dir2'), os.path.join(temp_dir, 'dir3'), os.path.join(temp_dir, 'non_dir4'), ] for p in dirs: os.makedirs(p) spill_dirs = os.path.pathsep.join([ os.path.join(temp_dir, 'select_dir'), os.path.join(temp_dir, 'dir*', 'subdir1'), ]) expected = sorted([dirs[0]] + [os.path.join(p, 'subdir1') for p in dirs[1:-1]]) self.assertEqual(expected, parse_spill_dirs(spill_dirs)) finally: shutil.rmtree(temp_dir)
def rm_spill_dirs(spill_dirs=None): import shutil spill_dirs = spill_dirs or [] if not isinstance(spill_dirs, list): spill_dirs = [spill_dirs] option_dirs = parse_spill_dirs(options.worker.spill_directory or '') spill_dirs = list(set(spill_dirs + option_dirs)) for d in spill_dirs: shutil.rmtree(d, ignore_errors=True)