def test_cleanup_dir():
def setup(file_names):
test_dir = tempfile.mkdtemp()
for file_name in file_names:
test_path = os.path.join(test_dir, file_name)
with open(test_path, 'w'):
pass
return test_dir
def tearDown(dir):
shutil.rmtree(dir)
# Test 1: Check if 'xgboost-model' is present after cleanup
model_name = "xgboost-model"
file_names = ['tmp1', 'tmp2', 'xgboost-model']
test_dir = setup(file_names)
train_utils.cleanup_dir(test_dir, model_name)
files = os.listdir(test_dir)
assert len(files) == 1
assert files[0] == model_name
tearDown(test_dir)
# Test 2: Check if directory is empty after cleanup
file_names = ['tmp1', 'tmp2']
test_dir = setup(file_names)
train_utils.cleanup_dir(test_dir, model_name)
files = os.listdir(test_dir)
assert len(files) == 0
tearDown(test_dir)
# Test 3: Check if directory is empty after cleanup
file_names = []
test_dir = setup(file_names)
train_utils.cleanup_dir(test_dir, model_name)
files = os.listdir(test_dir)
assert len(files) == 0
tearDown(test_dir)
def _cleanup_files(signo, frame):
if is_master:
train_utils.cleanup_dir(model_dir, MODEL_NAME)
_terminate()