From fdcbc29f26448f47201ec40fcf3b6405631c85d3 Mon Sep 17 00:00:00 2001
From: Tian Gao <gaogaotiantian@hotmail.com>
Date: Fri, 28 Feb 2025 13:15:55 -0500
Subject: gh-130660: Restore sys.ps1 and sys.ps2 after code.interact (#130661)

---
 Lib/test/test_code_module.py | 36 ++++++++++++++++++++++++++++++++----
 1 file changed, 32 insertions(+), 4 deletions(-)

(limited to 'Lib/test/test_code_module.py')

diff --git a/Lib/test/test_code_module.py b/Lib/test/test_code_module.py
index faa0b38f837..57fb130070b 100644
--- a/Lib/test/test_code_module.py
+++ b/Lib/test/test_code_module.py
@@ -39,19 +39,47 @@ class TestInteractiveConsole(unittest.TestCase, MockSys):
         self.mock_sys()
 
     def test_ps1(self):
-        self.infunc.side_effect = EOFError('Finished')
+        self.infunc.side_effect = [
+            "import code",
+            "code.sys.ps1",
+            EOFError('Finished')
+        ]
         self.console.interact()
-        self.assertEqual(self.sysmod.ps1, '>>> ')
+        output = ''.join(''.join(call[1]) for call in self.stdout.method_calls)
+        self.assertIn('>>> ', output)
+        self.assertNotHasAttr(self.sysmod, 'ps1')
+
+        self.infunc.side_effect = [
+            "import code",
+            "code.sys.ps1",
+            EOFError('Finished')
+        ]
         self.sysmod.ps1 = 'custom1> '
         self.console.interact()
+        output = ''.join(''.join(call[1]) for call in self.stdout.method_calls)
+        self.assertIn('custom1> ', output)
         self.assertEqual(self.sysmod.ps1, 'custom1> ')
 
     def test_ps2(self):
-        self.infunc.side_effect = EOFError('Finished')
+        self.infunc.side_effect = [
+            "import code",
+            "code.sys.ps2",
+            EOFError('Finished')
+        ]
         self.console.interact()
-        self.assertEqual(self.sysmod.ps2, '... ')
+        output = ''.join(''.join(call[1]) for call in self.stdout.method_calls)
+        self.assertIn('... ', output)
+        self.assertNotHasAttr(self.sysmod, 'ps2')
+
+        self.infunc.side_effect = [
+            "import code",
+            "code.sys.ps2",
+            EOFError('Finished')
+        ]
         self.sysmod.ps2 = 'custom2> '
         self.console.interact()
+        output = ''.join(''.join(call[1]) for call in self.stdout.method_calls)
+        self.assertIn('custom2> ', output)
         self.assertEqual(self.sysmod.ps2, 'custom2> ')
 
     def test_console_stderr(self):
-- 
cgit v1.2.3