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)