use unittest.addCleanup

sunmy2019 · sunmy2019 · commit ac72ac41c39a · 2023-05-30T18:41:01.000+08:00
diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
@@ -2320,6 +2320,7 @@ def test_variants(self):
         self.add_module_cleanup(name)
         with self.subTest(name):
             loaded = self.load(name)
+            self.addCleanup(loaded.module._clear_module_state)
 
             self.check_common(loaded)
             self.assertIsNot(loaded.snapshot.state_initialized, None)
@@ -2330,8 +2331,6 @@ def test_variants(self):
             self.assertIs(basic.look_up_self(), basic_lookedup)
             self.assertEqual(basic.initialized_count(), expected_init_count)
 
-            loaded.module._clear_module_state()
-
 
     def test_basic_reloaded(self):
         # m_copy is copied into the existing module object.
@@ -2391,6 +2390,10 @@ def test_with_reinit_reloaded(self):
                 loaded = self.load(name)
                 reloaded = self.re_load(name, loaded.module)
 
+                if name == f'{self.NAME}_with_state':
+                    self.addCleanup(loaded.module._clear_module_state)
+                    self.addCleanup(reloaded.module._clear_module_state)
+
                 self.check_common(loaded)
                 self.check_common(reloaded)
 
@@ -2416,10 +2419,6 @@ def test_with_reinit_reloaded(self):
 
                 self.assertIs(reloaded.snapshot.cached, reloaded.module)
 
-                if name == f'{self.NAME}_with_state':
-                    loaded.module._clear_module_state()
-                    reloaded.module._clear_module_state()
-
 
     # Currently, for every single-phrase init module loaded
     # in multiple interpreters, those interpreters share a
