From a99b9d911e0f8cb11b3436bdd8eb649b15d01a50 Mon Sep 17 00:00:00 2001
From: sunmy2019 <59365878+sunmy2019@users.noreply.github.com>
Date: Thu, 1 Jun 2023 05:12:10 +0800
Subject: [PATCH] gh-102251: Explicitly free state for test modules with state
 in test_import (#105085)

Co-authored-by: Erlend E. Aasland <erlend.aasland@protonmail.com>
---
 Lib/test/test_import/__init__.py | 14 ++++++++++----
 Modules/_testsinglephase.c       | 22 +++++++++++++++++++++-
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/Lib/test/test_import/__init__.py b/Lib/test/test_import/__init__.py
index e2384a08ecaa90..227c912bc8c322 100644
--- 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)
@@ -2379,14 +2380,19 @@ def test_with_reinit_reloaded(self):
         # Keep a reference around.
         basic = self.load(self.NAME)
 
-        for name in [
-            f'{self.NAME}_with_reinit',  # m_size == 0
-            f'{self.NAME}_with_state',  # m_size > 0
+        for name, has_state in [
+            (f'{self.NAME}_with_reinit', False),  # m_size == 0
+            (f'{self.NAME}_with_state', True),    # m_size > 0
         ]:
             self.add_module_cleanup(name)
-            with self.subTest(name):
+            with self.subTest(name=name, has_state=has_state):
                 loaded = self.load(name)
+                if has_state:
+                    self.addCleanup(loaded.module._clear_module_state)
+
                 reloaded = self.re_load(name, loaded.module)
+                if has_state:
+                    self.addCleanup(reloaded.module._clear_module_state)
 
                 self.check_common(loaded)
                 self.check_common(reloaded)
diff --git a/Modules/_testsinglephase.c b/Modules/_testsinglephase.c
index 8e6973f0b052cc..dca7abff89146e 100644
--- a/Modules/_testsinglephase.c
+++ b/Modules/_testsinglephase.c
@@ -248,6 +248,25 @@ basic__clear_globals(PyObject *self, PyObject *Py_UNUSED(ignored))
      basic__clear_globals_doc}
 
 
+PyDoc_STRVAR(basic__clear_module_state_doc, "_clear_module_state()\n\
+\n\
+Free the module state and set it to uninitialized.");
+
+static PyObject *
+basic__clear_module_state(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    module_state *state = get_module_state(self);
+    if (state != NULL) {
+        clear_state(state);
+    }
+    Py_RETURN_NONE;
+}
+
+#define _CLEAR_MODULE_STATE_METHODDEF \
+    {"_clear_module_state", basic__clear_module_state, METH_NOARGS, \
+     basic__clear_module_state_doc}
+
+
 /*********************************************/
 /* the _testsinglephase module (and aliases) */
 /*********************************************/
@@ -408,7 +427,7 @@ PyInit__testsinglephase_with_reinit(void)
 /* the _testsinglephase_with_state module */
 /******************************************/
 
-/* This ia less typical of legacy extensions in the wild:
+/* This is less typical of legacy extensions in the wild:
    - single-phase init  (same as _testsinglephase above)
    - has some module state
    - supports repeated initialization
@@ -424,6 +443,7 @@ static PyMethodDef TestMethods_WithState[] = {
     LOOK_UP_SELF_METHODDEF,
     SUM_METHODDEF,
     STATE_INITIALIZED_METHODDEF,
+    _CLEAR_MODULE_STATE_METHODDEF,
     {NULL, NULL}           /* sentinel */
 };