From 1426a99d31d012903550c98f038bce71137b1ee7 Mon Sep 17 00:00:00 2001 From: Benjamin Peterson Date: Sat, 3 Dec 2016 23:17:04 -0800 Subject: [PATCH] simplify cleanup of test_replace_parent_in_sys_modules (closes #28862) --- Lib/test/test_import.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/Lib/test/test_import.py b/Lib/test/test_import.py index 432500505e0..9a1fb678c61 100644 --- a/Lib/test/test_import.py +++ b/Lib/test/test_import.py @@ -400,20 +400,19 @@ class ImportTests(unittest.TestCase): def test_replace_parent_in_sys_modules(self): dir_name = os.path.abspath(TESTFN) os.mkdir(dir_name) - try: - pkg_dir = os.path.join(dir_name, 'sa') - os.mkdir(pkg_dir) - with open(os.path.join(pkg_dir, '__init__.py'), 'w') as init_file: - init_file.write("import v1") - with open(os.path.join(pkg_dir, 'v1.py'), 'w') as v1_file: - v1_file.write("import sys;" - "sys.modules['sa'] = sys.modules[__name__];" - "import sa") - sys.path.insert(0, dir_name) - # a segfault means the test failed! - import sa - finally: - rmtree(dir_name) + self.addCleanup(rmtree, dir_name) + pkg_dir = os.path.join(dir_name, 'sa') + os.mkdir(pkg_dir) + with open(os.path.join(pkg_dir, '__init__.py'), 'w') as init_file: + init_file.write("import v1") + with open(os.path.join(pkg_dir, 'v1.py'), 'w') as v1_file: + v1_file.write("import sys;" + "sys.modules['sa'] = sys.modules[__name__];" + "import sa") + sys.path.insert(0, dir_name) + self.addCleanup(sys.path.pop, 0) + # a segfault means the test failed! + import sa def test_fromlist_type(self): with self.assertRaises(TypeError) as cm: