refac: move run_tests to module scope

Author: willemt

pytest_asyncio_cooperative/plugin.py

@@ -122,16 +122,144 @@ def item_to_task(item):
         return test_wrapper(item)
 
 
-def _run_test_loop(tasks, session, run_tests):
+def get_coro(task):
+    if sys_version_info >= (3, 8):
+        return task.get_coro()
+    else:
+        return task._coro
+
+
+def cancel_task(task, now, item):
+    if sys_version_info >= (3, 9):
+        msg = "Test took too long ({:.2f} s)".format(now - item.enqueue_time)
+        task.cancel(msg=msg)
+    else:
+        task.cancel()
+
+
+async def run_tests(tasks, max_tasks: int, session, item_by_coro):
+    flakes_to_retry = []
+
+    sidelined_tasks = tasks[max_tasks:]
+    tasks = tasks[:max_tasks]
+
+    task_timeout = int(
+        session.config.getoption("--asyncio-task-timeout")
+        or session.config.getini("asyncio_task_timeout")
+    )
+
+    completed = []
+    cancelled = []
+    while tasks:
+        # Schedule all the coroutines
+        for i in range(len(tasks)):
+            if asyncio.iscoroutine(tasks[i]):
+                tasks[i] = asyncio.create_task(tasks[i])
+
+        # Mark when the task was started
+        earliest_enqueue_time = time.time()
+        for task in tasks:
+            if isinstance(task, asyncio.Task):
+                item = item_by_coro[get_coro(task)]
+            else:
+                item = item_by_coro[task]
+            if not hasattr(item, "enqueue_time"):
+                item.enqueue_time = time.time()
+            earliest_enqueue_time = min(item.enqueue_time, earliest_enqueue_time)
+
+        time_to_wait = (time.time() - earliest_enqueue_time) - task_timeout
+        done, pending = await asyncio.wait(
+            tasks,
+            return_when=asyncio.FIRST_COMPLETED,
+            timeout=min(30, int(time_to_wait)),
+        )
+
+        # Cancel tasks that have taken too long
+        tasks = []
+        for task in pending:
+            now = time.time()
+            item = item_by_coro[get_coro(task)]
+            if task not in cancelled and task_timeout < now - item.enqueue_time:
+                cancel_task(task, now, item)
+                cancelled.append(task)
+            tasks.append(task)
+
+        for result in done:
+            item = item_by_coro[get_coro(result)]
+
+            # Flakey tests will be run again if they failed
+            # TODO: add retry count
+            if item._flakey:
+                try:
+                    result.result()
+                except:
+                    item._flakey = None
+                    new_task = item_to_task(item)
+                    flakes_to_retry.append(new_task)
+                    item_by_coro[new_task] = item
+                    continue
+
+            # We need to change .runtest to a synchronous function for pytest
+            # however, if it is called again by retry libraries we need to rerun
+            # the test instead of retuning the previous result
+            def wrap_in_sync():
+                def sync_wrapper():
+                    new_task = item_to_task(item)
+
+                    # We use a new thread because we can't block for an async function
+                    # in the same thread as the current running event loop, nor
+                    # we can nest event loops
+                    result = None
+
+                    def run_in_thread():
+                        nonlocal result
+                        try:
+                            result = asyncio.run(new_task)
+                        except Exception as e:
+                            result = e
+
+                    thread = threading.Thread(target=run_in_thread)
+                    thread.start()
+                    thread.join()
+
+                    if isinstance(result, Exception):
+                        raise result  # type: ignore
+
+                    return result
+
+                item.runtest = sync_wrapper
+
+                return result.result()
+
+            item.runtest = wrap_in_sync
+
+            item.ihook.pytest_runtest_protocol(item=item, nextitem=None)
+
+            # Hack: See rewrite comment below
+            # pytest_runttest_protocl will disable the rewrite assertion
+            # so we renable it here
+            activate_assert_rewrite(item)
+
+            completed.append(result)
+
+        if sidelined_tasks:
+            if len(tasks) < max_tasks:
+                tasks.append(sidelined_tasks.pop(0))
+
+    return flakes_to_retry
+
+
+def _run_test_loop(tasks, session, item_by_coro):
     max_tasks = int(
         session.config.getoption("--max-asyncio-tasks")
         or session.config.getini("max_asyncio_tasks")
     )
 
     loop = asyncio.new_event_loop()
     try:
-        task = run_tests(tasks, int(max_tasks))
-        loop.run_until_complete(task)
+        return loop.run_until_complete(
+            run_tests(tasks, int(max_tasks), session, item_by_coro)
+        )
     finally:
         loop.close()
 
@@ -155,8 +283,6 @@ def pytest_runtestloop(session):
 
     session.wrapped_fixtures = {}
 
-    flakes_to_retry = []
-
     # Collect our coroutines
     regular_items = []
     item_by_coro = {}
@@ -180,127 +306,6 @@ def pytest_runtestloop(session):
         else:
             regular_items.append(item)
 
-    def get_coro(task):
-        if sys_version_info >= (3, 8):
-            return task.get_coro()
-        else:
-            return task._coro
-
-    async def run_tests(tasks, max_tasks: int):
-        sidelined_tasks = tasks[max_tasks:]
-        tasks = tasks[:max_tasks]
-
-        task_timeout = int(
-            session.config.getoption("--asyncio-task-timeout")
-            or session.config.getini("asyncio_task_timeout")
-        )
-
-        completed = []
-        cancelled = []
-        while tasks:
-            # Schedule all the coroutines
-            for i in range(len(tasks)):
-                if asyncio.iscoroutine(tasks[i]):
-                    tasks[i] = asyncio.create_task(tasks[i])
-
-            # Mark when the task was started
-            earliest_enqueue_time = time.time()
-            for task in tasks:
-                if isinstance(task, asyncio.Task):
-                    item = item_by_coro[get_coro(task)]
-                else:
-                    item = item_by_coro[task]
-                if not hasattr(item, "enqueue_time"):
-                    item.enqueue_time = time.time()
-                earliest_enqueue_time = min(item.enqueue_time, earliest_enqueue_time)
-
-            time_to_wait = (time.time() - earliest_enqueue_time) - task_timeout
-            done, pending = await asyncio.wait(
-                tasks,
-                return_when=asyncio.FIRST_COMPLETED,
-                timeout=min(30, int(time_to_wait)),
-            )
-
-            # Cancel tasks that have taken too long
-            tasks = []
-            for task in pending:
-                now = time.time()
-                item = item_by_coro[get_coro(task)]
-                if task not in cancelled and task_timeout < now - item.enqueue_time:
-                    if sys_version_info >= (3, 9):
-                        msg = "Test took too long ({:.2f} s)".format(
-                            now - item.enqueue_time
-                        )
-                        task.cancel(msg=msg)
-                    else:
-                        task.cancel()
-                    cancelled.append(task)
-                tasks.append(task)
-
-            for result in done:
-                item = item_by_coro[get_coro(result)]
-
-                # Flakey tests will be run again if they failed
-                # TODO: add retry count
-                if item._flakey:
-                    try:
-                        result.result()
-                    except:
-                        item._flakey = None
-                        new_task = item_to_task(item)
-                        flakes_to_retry.append(new_task)
-                        item_by_coro[new_task] = item
-                        continue
-
-                # We need to change .runtest to a synchronous function for pytest
-                # however, if it is called again by retry libraries we need to rerun
-                # the test instead of retuning the previous result
-                def wrap_in_sync():
-                    def sync_wrapper():
-                        new_task = item_to_task(item)
-
-                        # We use a new thread because we can't block for an async function
-                        # in the same thread as the current running event loop, nor
-                        # we can nest event loops
-                        result = None
-
-                        def run_in_thread():
-                            nonlocal result
-                            try:
-                                result = asyncio.run(new_task)
-                            except Exception as e:
-                                result = e
-
-                        thread = threading.Thread(target=run_in_thread)
-                        thread.start()
-                        thread.join()
-
-                        if isinstance(result, Exception):
-                            raise result  # type: ignore
-
-                        return result
-
-                    item.runtest = sync_wrapper
-
-                    return result.result()
-
-                item.runtest = wrap_in_sync
-
-                item.ihook.pytest_runtest_protocol(item=item, nextitem=None)
-
-                # Hack: See rewrite comment below
-                # pytest_runttest_protocl will disable the rewrite assertion
-                # so we renable it here
-                activate_assert_rewrite(item)
-
-                completed.append(result)
-
-            if sidelined_tasks:
-                if len(tasks) < max_tasks:
-                    tasks.append(sidelined_tasks.pop(0))
-
-        return completed
-
     # Do assert rewrite
     # Hack: pytest's implementation sets up assert rewriting as a shared
     # resource. This causes a race condition between async tests. Therefore we
@@ -313,11 +318,11 @@ def run_in_thread():
         return
 
     # Run the tests using cooperative multitasking
-    _run_test_loop(tasks, session, run_tests)
+    flakes_to_retry = _run_test_loop(tasks, session, item_by_coro)
 
     # Run failed flakey tests
     if flakes_to_retry:
-        _run_test_loop(flakes_to_retry, session, run_tests)
+        _run_test_loop(flakes_to_retry, session, item_by_coro)
 
     # Run synchronous tests
     session.items = regular_items