from setUp and tearDown to the custom context manager

comtypes/test/test_findgendir.py

@@ -1,8 +1,10 @@
+import contextlib
 import importlib
 import os
 import sys
 import tempfile
 import types
+from typing import Iterator
 import unittest
 from unittest import mock
 
@@ -14,34 +16,32 @@
 IMGBASE = os.path.splitext(os.path.basename(sys.executable))[0]
 
 
-class Test(unittest.TestCase):
-    """Test the comtypes.client._find_gen_dir() function in several
-    simulated environments.
-    """
-
-    def setUp(self):
-        # save the original comtypes.gen modules and create a
-        # substitute with an empty __path__.
-        self.orig_comtypesgen = sys.modules["comtypes.gen"]
+@contextlib.contextmanager
+def patch_empty_module_to_comtypes_gen() -> Iterator[types.ModuleType]:
+    with mock.patch.dict(sys.modules):
         del sys.modules["comtypes.gen"]
-        del comtypes.gen
         mod = sys.modules["comtypes.gen"] = types.ModuleType("comtypes.gen")
         mod.__path__ = []
-        comtypes.gen = mod
+        with mock.patch.object(comtypes, "gen", mod):
+            try:
+                yield mod
+            finally:
+                importlib.reload(comtypes.gen)
 
-    def tearDown(self):
-        # restore the original comtypes.gen module
-        comtypes.gen = self.orig_comtypesgen
-        sys.modules["comtypes.gen"] = self.orig_comtypesgen
-        importlib.reload(comtypes.gen)
+
+class Test(unittest.TestCase):
+    """Test the comtypes.client._find_gen_dir() function in several
+    simulated environments.
+    """
 
     def test_script(self):
         # %APPDATA%\Python\Python25\comtypes_cache
         ma, mi = sys.version_info[:2]
         cache = rf"$APPDATA\Python\Python{ma:d}{mi:d}\comtypes_cache"
         path = os.path.expandvars(cache)
-        gen_dir = comtypes.client._find_gen_dir()
-        self.assertEqual(path, gen_dir)
+        with patch_empty_module_to_comtypes_gen():
+            gen_dir = comtypes.client._find_gen_dir()
+            self.assertEqual(path, gen_dir)
 
     # patch py2exe-attributes to `sys` modules
     @mock.patch.object(sys, "frozen", "dll", create=True)
@@ -52,8 +52,9 @@ def test_frozen_dll(self):
         ma, mi = sys.version_info[:2]
         cache = rf"comtypes_cache\{IMGBASE}{ma:d}{mi:d}-{ma:d}{mi:d}"
         path = os.path.join(tempfile.gettempdir(), cache)
-        gen_dir = comtypes.client._find_gen_dir()
-        self.assertEqual(path, gen_dir)
+        with patch_empty_module_to_comtypes_gen():
+            gen_dir = comtypes.client._find_gen_dir()
+            self.assertEqual(path, gen_dir)
 
     # patch py2exe-attributes to `sys` modules
     @mock.patch.object(sys, "frozen", "console_exe", create=True)
@@ -62,8 +63,9 @@ def test_frozen_console_exe(self):
         ma, mi = sys.version_info[:2]
         cache = rf"comtypes_cache\{IMGBASE}-{ma:d}{mi:d}"
         path = os.path.join(tempfile.gettempdir(), cache)
-        gen_dir = comtypes.client._find_gen_dir()
-        self.assertEqual(path, gen_dir)
+        with patch_empty_module_to_comtypes_gen():
+            gen_dir = comtypes.client._find_gen_dir()
+            self.assertEqual(path, gen_dir)
 
     # patch py2exe-attributes to `sys` modules
     @mock.patch.object(sys, "frozen", "windows_exe", create=True)
@@ -72,8 +74,9 @@ def test_frozen_windows_exe(self):
         ma, mi = sys.version_info[:2]
         cache = rf"comtypes_cache\{IMGBASE}-{ma:d}{mi:d}"
         path = os.path.join(tempfile.gettempdir(), cache)
-        gen_dir = comtypes.client._find_gen_dir()
-        self.assertEqual(path, gen_dir)
+        with patch_empty_module_to_comtypes_gen():
+            gen_dir = comtypes.client._find_gen_dir()
+            self.assertEqual(path, gen_dir)
 
 
 def main():