Remove some more boilerplate from the actual tests in test_pdb.

1 files changed, 34 insertions, 27 deletions

diff --git a/Lib/test/test_pdb.py b/Lib/test/test_pdb.py
index 68a7702..a15538c 100644
--- a/Lib/test/test_pdb.py
+++ b/Lib/test/test_pdb.py
@@ -12,23 +12,33 @@ from test import test_support
 from test_doctest import _FakeInput
 
 
+class PdbTestInput(object):
+    """Context manager that makes testing Pdb in doctests easier."""
+
+    def __init__(self, input):
+        self.input = input
+
+    def __enter__(self):
+        self.real_stdin = sys.stdin
+        sys.stdin = _FakeInput(self.input)
+
+    def __exit__(self, *exc):
+        sys.stdin = self.real_stdin
+
+
 def test_pdb_skip_modules():
     """This illustrates the simple case of module skipping.
 
     >>> def skip_module():
     ...     import string
-    ...     import pdb;pdb.Pdb(skip=['string*']).set_trace()
+    ...     import pdb; pdb.Pdb(skip=['string*']).set_trace()
     ...     string.lower('FOO')
-    >>> real_stdin = sys.stdin
-    >>> sys.stdin = _FakeInput([
-    ...    'step',
-    ...    'continue',
-    ...    ])
 
-    >>> try:
+    >>> with PdbTestInput([
+    ...     'step',
+    ...     'continue'
+    ... ]):
     ...     skip_module()
-    ... finally:
-    ...     sys.stdin = real_stdin
     > <doctest test.test_pdb.test_pdb_skip_modules[0]>(4)skip_module()
     -> string.lower('FOO')
     (Pdb) step
@@ -36,7 +46,7 @@ def test_pdb_skip_modules():
     > <doctest test.test_pdb.test_pdb_skip_modules[0]>(4)skip_module()->None
     -> string.lower('FOO')
     (Pdb) continue
-"""
+    """
 
 
 # Module for testing skipping of module that makes a callback
@@ -50,22 +60,19 @@ def test_pdb_skip_modules_with_callback():
     >>> def skip_module():
     ...     def callback():
     ...         return None
-    ...     import pdb;pdb.Pdb(skip=['module_to_skip*']).set_trace()
+    ...     import pdb; pdb.Pdb(skip=['module_to_skip*']).set_trace()
     ...     mod.foo_pony(callback)
-    >>> real_stdin = sys.stdin
-    >>> sys.stdin = _FakeInput([
-    ...    'step',
-    ...    'step',
-    ...    'step',
-    ...    'step',
-    ...    'step',
-    ...    'continue',
-    ...    ])
-
-    >>> try:
+
+    >>> with PdbTestInput([
+    ...     'step',
+    ...     'step',
+    ...     'step',
+    ...     'step',
+    ...     'step',
+    ...     'continue',
+    ... ]):
     ...     skip_module()
-    ... finally:
-    ...     sys.stdin = real_stdin
+    ...     pass  # provides something to "step" to
     > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(5)skip_module()
     -> mod.foo_pony(callback)
     (Pdb) step
@@ -84,10 +91,10 @@ def test_pdb_skip_modules_with_callback():
     > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[0]>(5)skip_module()->None
     -> mod.foo_pony(callback)
     (Pdb) step
-    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[3]>(4)<module>()
-    -> sys.stdin = real_stdin
+    > <doctest test.test_pdb.test_pdb_skip_modules_with_callback[1]>(10)<module>()
+    -> pass  # provides something to "step" to
     (Pdb) continue
-"""
+    """
 
 
 def test_main():