test for the Stop(Async)Iteration case

graingert · graingert · commit adf2bf3cc554 · 2022-08-11T14:52:52.000+01:00
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
@@ -104,15 +104,41 @@ def f():
         self.assertEqual(frames[0].line, '1/0')
 
         # Repeat with RuntimeError (which goes through a different code path)
+        class RuntimeErrorSubclass(RuntimeError):
+            pass
+
         try:
             with f():
-                raise NotImplementedError(42)
-        except NotImplementedError as e:
+                raise RuntimeErrorSubclass(42)
+        except RuntimeErrorSubclass as e:
             frames = traceback.extract_tb(e.__traceback__)
 
         self.assertEqual(len(frames), 1)
         self.assertEqual(frames[0].name, 'test_contextmanager_traceback')
-        self.assertEqual(frames[0].line, 'raise NotImplementedError(42)')
+        self.assertEqual(frames[0].line, 'raise RuntimeErrorSubclass(42)')
+
+        class StopIterationSubclass(StopIteration):
+            pass
+
+        for stop_exc in (
+            StopIteration('spam'),
+            StopIterationSubclass('spam'),
+        ):
+            with self.subTest(type=type(stop_exc)):
+                try:
+                    with f():
+                        raise stop_exc
+                except type(stop_exc) as e:
+                    self.assertIs(e, stop_exc)
+                    frames = traceback.extract_tb(e.__traceback__)
+                else:
+                    self.fail(f'{stop_exc} was suppressed')
+
+                self.assertEqual(len(frames), 2)
+                self.assertEqual(frames[0].name, 'f')
+                self.assertEqual(frames[0].line, 'yield')
+                self.assertEqual(frames[1].name, 'test_contextmanager_traceback')
+                self.assertEqual(frames[1].line, 'raise stop_exc')
 
     def test_contextmanager_no_reraise(self):
         @contextmanager
diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py
@@ -143,15 +143,46 @@ async def f():
         self.assertEqual(frames[0].line, '1/0')
 
         # Repeat with RuntimeError (which goes through a different code path)
+        class RuntimeErrorSubclass(RuntimeError):
+            pass
+
         try:
             async with f():
-                raise NotImplementedError(42)
-        except NotImplementedError as e:
+                raise RuntimeErrorSubclass(42)
+        except RuntimeErrorSubclass as e:
             frames = traceback.extract_tb(e.__traceback__)
 
         self.assertEqual(len(frames), 1)
         self.assertEqual(frames[0].name, 'test_contextmanager_traceback')
-        self.assertEqual(frames[0].line, 'raise NotImplementedError(42)')
+        self.assertEqual(frames[0].line, 'raise RuntimeErrorSubclass(42)')
+
+        class StopIterationSubclass(StopIteration):
+            pass
+
+        class StopAsyncIterationSubclass(StopAsyncIteration):
+            pass
+
+        for stop_exc in (
+            StopIteration('spam'),
+            StopAsyncIteration('ham'),
+            StopIterationSubclass('spam'),
+            StopAsyncIterationSubclass('spam')
+        ):
+            with self.subTest(type=type(stop_exc)):
+                try:
+                    async with f():
+                        raise stop_exc
+                except type(stop_exc) as e:
+                    self.assertIs(e, stop_exc)
+                    frames = traceback.extract_tb(e.__traceback__)
+                else:
+                    self.fail(f'{stop_exc} was suppressed')
+
+                self.assertEqual(len(frames), 2)
+                self.assertEqual(frames[0].name, 'f')
+                self.assertEqual(frames[0].line, 'yield')
+                self.assertEqual(frames[1].name, 'test_contextmanager_traceback')
+                self.assertEqual(frames[1].line, 'raise stop_exc')
 
     @_async_test
     async def test_contextmanager_no_reraise(self):
