Spaces:
Running
Running
File size: 1,042 Bytes
6a86ad5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
"""
This module adds context manager for temporary files generated by the tests.
"""
import shutil
import os
class TmpFileManager:
"""
A class to track record of every temporary files created by the tests.
"""
tmp_files = set('')
tmp_folders = set('')
@classmethod
def tmp_file(cls, name=''):
cls.tmp_files.add(name)
return name
@classmethod
def tmp_folder(cls, name=''):
cls.tmp_folders.add(name)
return name
@classmethod
def cleanup(cls):
while cls.tmp_files:
file = cls.tmp_files.pop()
if os.path.isfile(file):
os.remove(file)
while cls.tmp_folders:
folder = cls.tmp_folders.pop()
shutil.rmtree(folder)
def cleanup_tmp_files(test_func):
"""
A decorator to help test codes remove temporary files after the tests.
"""
def wrapper_function():
try:
test_func()
finally:
TmpFileManager.cleanup()
return wrapper_function
|