This recipe helps to manage a temporary file used in a test. It returns an open file so as to write the test data, and upon exiting the 'with' statement, makes sure that the created file is deleted.
1 2 3 4 5 6 7 8 9 10 11 12
from contextlib import contextmanager import os @contextmanager def test_file(path): try: open_file = open(path, 'w') yield open_file finally: open_file.close() if os.path.exists(path): os.unlink(path)
unittest.TestCase.setUp/tearDown are not always the right level of granularity for tests. Sometimes what needs to be written to a test file varies widely between tests. It can be convenient to have unique code per test for writing out to a test file.
This context manager is for those situations. Just write to the open file, close it, and then use the file path or open a new file object within the 'with' statement. You are guaranteed that the file will be deleted at the end of your code.