asyncio: Refactor tests: add a base TestCase class

1 files changed, 18 insertions, 0 deletions

diff --git a/Lib/asyncio/test_utils.py b/Lib/asyncio/test_utils.py
index 1062bae132..d9c7ae2d11 100644
--- a/Lib/asyncio/test_utils.py
+++ b/Lib/asyncio/test_utils.py
@@ -11,6 +11,7 @@ import sys
 import tempfile
 import threading
 import time
+import unittest
 from unittest import mock
 
 from http.server import HTTPServer
@@ -379,3 +380,20 @@ def get_function_source(func):
     if source is None:
         raise ValueError("unable to get the source of %r" % (func,))
     return source
+
+
+class TestCase(unittest.TestCase):
+    def set_event_loop(self, loop, *, cleanup=True):
+        assert loop is not None
+        # ensure that the event loop is passed explicitly in asyncio
+        events.set_event_loop(None)
+        if cleanup:
+            self.addCleanup(loop.close)
+
+    def new_test_loop(self, gen=None):
+        loop = TestLoop(gen)
+        self.set_event_loop(loop)
+        return loop
+
+    def tearDown(self):
+        events.set_event_loop(None)