Implement concurrent limits based on item marks

This change enables users to limit individual test concurrency based
on mark configuration. For example, lets say that there are 10 tests,
and 8 of them can be executed concurrently, but 2 of them use some
shared resource and can't be executed concurrenctly. By using the new
option `--max-asyncio-tasks-by-mark`, its now possible to configure
the task scheduler to execute the 2 tasks in series, but the remaining
tasks 8 tasks would still be scheduled concurrently.

Author: arg0d

.gitignore

@@ -0,0 +1 @@
+__pycache__

pytest_asyncio_cooperative/fixtures.py

@@ -23,7 +23,7 @@ def _get_fixture(item, arg_name, fixture=None):
         if fixture:
             try:
                 item._request.param = item._pyfuncitem.callspec.params[fixture.argname]
-            except (AttributeError, KeyError) :
+            except (AttributeError, KeyError):
                 pass
 
         return item._request
@@ -62,16 +62,19 @@ async def fill_fixtures(item):
 
         if fixture.scope not in ["function", "module", "session"]:
             raise Exception(f"{fixture.scope} scope not supported")
-        
+
         fixtures.append(fixture)
         are_autouse.append(is_autouse)
-    
+
     # Fill fixtures concurrently
     fill_results = await asyncio.gather(
-        *(fill_fixture_fixtures(item._fixtureinfo, fixture, item) for fixture in fixtures)
+        *(
+            fill_fixture_fixtures(item._fixtureinfo, fixture, item)
+            for fixture in fixtures
+        )
     )
 
-    for ((value, extra_teardowns), is_autouse) in zip(fill_results, are_autouse):
+    for (value, extra_teardowns), is_autouse in zip(fill_results, are_autouse):
         teardowns.extend(extra_teardowns)
 
         if not is_autouse:
@@ -355,7 +358,9 @@ async def fill_fixture_fixtures(_fixtureinfo, fixture, item):
     ):
         return await _make_coroutine_fixture(_fixtureinfo, fixture, item)
 
-    elif inspect.isgeneratorfunction(fixture.func) or isinstance(fixture.func, CachedGen):
+    elif inspect.isgeneratorfunction(fixture.func) or isinstance(
+        fixture.func, CachedGen
+    ):
         return await _make_regular_generator_fixture(_fixtureinfo, fixture, item)
 
     elif inspect.isfunction(fixture.func):

pytest_asyncio_cooperative/plugin.py

@@ -27,6 +27,18 @@ def pytest_addoption(parser):
         default=100,
     )
 
+    parser.addoption(
+        "--max-asyncio-tasks-by-mark",
+        action="store",
+        default=None,
+        help="asyncio: maximum number of tasks to run concurrently for each mark (mark=int pairs)",
+    )
+    parser.addini(
+        "max_asyncio_tasks_by_mark",
+        "asyncio: asyncio: maximum number of tasks to run concurrently for each mark (mark=int pairs)",
+        default=None,
+    )
+
     parser.addoption(
         "--asyncio-task-timeout",
         action="store",
@@ -171,15 +183,77 @@ def item_to_task(item):
         raise NotCoroutine
 
 
+def parse_max_tasks_by_mark(config, debug_prefix):
+    if not config:
+        return {}
+
+    result = {}
+    pairs = config.split(",")
+    for pair in pairs:
+        columns = pair.split("=")
+        if len(columns) > 2:
+            assert False, f"`{debug_prefix}`: too many `=` in `{pair}`"
+        try:
+            max_tasks = int(columns[1])
+        except ValueError:
+            assert False, f"`{debug_prefix}`: expected integer in `{pair}`"
+
+        mark = columns[0]
+        result[mark] = max_tasks
+
+    return result
+
+
+def item_would_exceed_max_marks(item, active_marks, max_tasks_by_mark):
+    for mark in item.own_markers:
+        if (
+            mark.name in active_marks
+            and active_marks[mark.name] >= max_tasks_by_mark[mark.name]
+        ):
+            return True
+    return False
+
+
+def item_add_active_marks(item, delta, active_marks):
+    for mark in item.own_markers:
+        if mark.name in active_marks:
+            active_marks[mark.name] += delta
+
+
+def get_coro(task):
+    if sys_version_info >= (3, 8):
+        return task.get_coro()
+    else:
+        return task._coro
+
+
+def get_item_by_coro(task, item_by_coro):
+    if isinstance(task, asyncio.Task):
+        return item_by_coro[get_coro(task)]
+    else:
+        return item_by_coro[task]
+
+
 def _run_test_loop(tasks, session, run_tests):
     max_tasks = int(
         session.config.getoption("--max-asyncio-tasks")
         or session.config.getini("max_asyncio_tasks")
     )
 
+    if session.config.getoption("--max-asyncio-tasks-by-mark"):
+        max_tasks_by_mark = parse_max_tasks_by_mark(
+            session.config.getoption("--max-asyncio-tasks-by-mark"),
+            "--max-asyncio-tasks-by-mark",
+        )
+    else:
+        max_tasks_by_mark = parse_max_tasks_by_mark(
+            session.config.getini("max_asyncio_tasks_by_mark"),
+            "max_asyncio_tasks_by_mark",
+        )
+
     loop = asyncio.new_event_loop()
     try:
-        task = run_tests(tasks, int(max_tasks))
+        task = run_tests(tasks, int(max_tasks), max_tasks_by_mark)
         loop.run_until_complete(task)
     finally:
         loop.close()
@@ -234,15 +308,35 @@ 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, max_tasks_by_mark):
+        sidelined_tasks = tasks
+        tasks = []
+        active_marks = {mark: 0 for mark, max_tasks in max_tasks_by_mark.items()}
+
+        def enqueue_tasks():
+            while len(tasks) < max_tasks:
+                remove_index = None
+
+                for index, task in enumerate(sidelined_tasks):
+                    item = get_item_by_coro(task, item_by_coro)
+                    if not item_would_exceed_max_marks(
+                        item, active_marks, max_tasks_by_mark
+                    ):
+                        item_add_active_marks(item, 1, active_marks)
+                        tasks.append(task)
+                        remove_index = index
+                        break
+
+                if remove_index == None:
+                    # No available tasks were found, give up control.
+                    break
+
+                # Removing element from start/middle of the list is actually
+                # quite fast compared to iterating to the end of the list in
+                # the above loop.
+                sidelined_tasks.pop(remove_index)
 
-    async def run_tests(tasks, max_tasks: int):
-        sidelined_tasks = tasks[max_tasks:]
-        tasks = tasks[:max_tasks]
+        enqueue_tasks()
 
         task_timeout = int(
             session.config.getoption("--asyncio-task-timeout")
@@ -260,17 +354,16 @@ async def run_tests(tasks, max_tasks: int):
             # 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]
+                item = get_item_by_coro(task, item_by_coro)
                 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))
+                tasks,
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=min(30, int(time_to_wait)),
             )
 
             # Cancel tasks that have taken too long
@@ -291,6 +384,7 @@ async def run_tests(tasks, max_tasks: int):
 
             for result in done:
                 item = item_by_coro[get_coro(result)]
+                item_add_active_marks(item, -1, active_marks)
 
                 # Flakey tests will be run again if they failed
                 # TODO: add retry count
@@ -315,9 +409,7 @@ async def run_tests(tasks, max_tasks: int):
 
                 completed.append(result)
 
-            if sidelined_tasks:
-                if len(tasks) < max_tasks:
-                    tasks.append(sidelined_tasks.pop(0))
+            enqueue_tasks()
 
         return completed
 

tests/test_class_based.py

@@ -1,5 +1,6 @@
 def test_class_based_tests(testdir):
-    testdir.makepyfile("""
+    testdir.makepyfile(
+        """
         import pytest
 
 
@@ -8,15 +9,17 @@ class TestSuite:
             async def test_cooperative(self):
                 assert True
 
-    """)
+    """
+    )
 
     result = testdir.runpytest()
 
     result.assert_outcomes(passed=1)
 
 
 def test_class_based_tests_with_fixture(testdir):
-    testdir.makepyfile("""
+    testdir.makepyfile(
+        """
         import pytest
 
 
@@ -29,7 +32,8 @@ async def test_fixture(self):
             async def test_cooperative(self, test_fixture):
                 assert test_fixture == "test_fixture"
 
-    """)
+    """
+    )
 
     result = testdir.runpytest()
 

tests/test_fail.py

@@ -30,13 +30,15 @@ def test_a():
     result.assert_outcomes(failed=1)
 
 
-@pytest.mark.parametrize("dur1, dur2, expectedfails, expectedpasses", [
-    (1.1, 2, 2, 0),
-    (2, 2, 2, 0),
-])
+@pytest.mark.parametrize(
+    "dur1, dur2, expectedfails, expectedpasses",
+    [
+        (1.1, 2, 2, 0),
+        (2, 2, 2, 0),
+    ],
+)
 def test_function_takes_too_long(testdir, dur1, dur2, expectedfails, expectedpasses):
-    testdir.makeconftest(
-        """""")
+    testdir.makeconftest("""""")
 
     testdir.makepyfile(
         """
@@ -51,7 +53,9 @@ async def test_a():
         @pytest.mark.asyncio_cooperative
         async def test_b():
             await asyncio.sleep({})
-    """.format(dur1, dur2)
+    """.format(
+            dur1, dur2
+        )
     )
 
     result = testdir.runpytest("--asyncio-task-timeout", "1")

tests/test_fixture.py

@@ -364,7 +364,9 @@ async def test_ordering(ghi, _def, abc):
     result.assert_outcomes(passed=1)
 
 
-def test_ordering_of_fixtures_based_off_function_arguments_with_session_fixture(testdir):
+def test_ordering_of_fixtures_based_off_function_arguments_with_session_fixture(
+    testdir,
+):
     testdir.makepyfile(
         """
         import asyncio