Repository: pytest-dev/pytest-xdist
Pull request #733 (2021-11-28 15:11:01)

Create new dist option 'loadgroup'
### Checklist
- [x] Make sure to include reasonable tests for your change if necessary
- [x] Add change log

### Intro
There are currently several options for distributing tests,
but there is still no suitable option for the following cases:

### Case 1
In this case, it is efficient to divide all tests into different sessions.
```python
@pytest.mark.parametrize('param', [A, B, C, D])
def test_something_heavy(param):
    do_something_heavy_test
```

### Case 2
In this case, it is efficient to run all tests in the same session.
```python
def test_something_light_1(heavy_fixture_cannot_filelock):
    do_something_light_test

def test_something_light_2(heavy_fixture_cannot_filelock):
    do_something_light_test

def test_something_light_3(heavy_fixture_cannot_filelock):
    do_something_light_test
```

### Limit
If you use the loadscope option, all tests in Case 1 are performed in same session,
If the load option is used, all tests of Case 2 may be performed in different sessions.

### Suggestion
Use the following xgroup mark and specify the name through the parameter.
Then, tests with the same name are executed in the same session.
```python
@pytest.mark.xgroup(name="same_session")
def test_something_light_1(heavy_fixture_cannot_filelock):
    do_something_light_test

@pytest.mark.xgroup(name="same_session")
def test_something_light_2(heavy_fixture_cannot_filelock):
    do_something_light_test

@pytest.mark.xgroup(name="same_session")
def test_something_light_3(heavy_fixture_cannot_filelock):
    do_something_light_test
```



Patch:
diff --git a/README.rst b/README.rst
index 5768f7c..c40c9f3 100644
--- a/README.rst
+++ b/README.rst
@@ -96,6 +96,10 @@ distribution algorithm this with the ``--dist`` option. It takes these values:
   distributed to available workers as whole units. This guarantees that all
   tests in a file run in the same worker.
 
+* ``--dist loadgroup``: Tests are grouped by xdist_group mark. Groups are
+  distributed to available workers as whole units. This guarantees that all
+  tests with same xdist_group name run in the same worker.
+
 Making session-scoped fixtures execute only once
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -414,3 +418,21 @@ where the configuration file was found.
 .. _`pytest-xdist`: http://pypi.python.org/pypi/pytest-xdist
 .. _`pytest-xdist repository`: https://github.com/pytest-dev/pytest-xdist
 .. _`pytest`: http://pytest.org
+
+Groups tests by xdist_group mark
+---------------------------------
+
+*New in version 2.4.*
+
+Two or more tests belonging to different classes or modules can be executed in same worker through the xdist_group marker:
+
+.. code-block:: python
+
+    @pytest.mark.xdist_group(name="group1")
+    def test1():
+        pass
+
+    class TestA:
+        @pytest.mark.xdist_group("group1")
+        def test2():
+            pass
diff --git a/changelog/733.feature.rst b/changelog/733.feature.rst
new file mode 100644
index 0000000..28163e7
--- /dev/null
+++ b/changelog/733.feature.rst
@@ -0,0 +1,1 @@
+New ``--dist=loadgroup`` option, which ensures all tests marked with ``@pytest.mark.xdist_group`` run in the same session/worker. Other tests run distributed as in ``--dist=load``.
diff --git a/src/xdist/dsession.py b/src/xdist/dsession.py
index 12539a0..eb5ae75 100644
--- a/src/xdist/dsession.py
+++ b/src/xdist/dsession.py
@@ -7,6 +7,7 @@ from xdist.scheduler import (
     LoadScheduling,
     LoadScopeScheduling,
     LoadFileScheduling,
+    LoadGroupScheduling,
 )
 
 
@@ -100,6 +101,7 @@ class DSession:
             "load": LoadScheduling,
             "loadscope": LoadScopeScheduling,
             "loadfile": LoadFileScheduling,
+            "loadgroup": LoadGroupScheduling,
         }
         return schedulers[dist](config, log)
 
diff --git a/src/xdist/plugin.py b/src/xdist/plugin.py
index 17be6d0..85f76e8 100644
--- a/src/xdist/plugin.py
+++ b/src/xdist/plugin.py
@@ -86,7 +86,7 @@ def pytest_addoption(parser):
         "--dist",
         metavar="distmode",
         action="store",
-        choices=["each", "load", "loadscope", "loadfile", "no"],
+        choices=["each", "load", "loadscope", "loadfile", "loadgroup", "no"],
         dest="dist",
         default="no",
         help=(
@@ -98,6 +98,7 @@ def pytest_addoption(parser):
             " the same scope to any available environment.\n\n"
             "loadfile: load balance by sending test grouped by file"
             " to any available environment.\n\n"
+            "loadgroup: like load, but sends tests marked with 'xdist_group' to the same worker.\n\n"
             "(default) no: run tests inprocess, don't distribute."
         ),
     )
@@ -204,6 +205,12 @@ def pytest_configure(config):
         config.issue_config_time_warning(warning, 2)
         config.option.forked = True
 
+    config_line = (
+        "xdist_group: specify group for tests should run in same session."
+        "in relation to one another. " + "Provided by pytest-xdist."
+    )
+    config.addinivalue_line("markers", config_line)
+
 
 @pytest.hookimpl(tryfirst=True)
 def pytest_cmdline_main(config):
diff --git a/src/xdist/remote.py b/src/xdist/remote.py
index 914040d..160b042 100644
--- a/src/xdist/remote.py
+++ b/src/xdist/remote.py
@@ -116,6 +116,20 @@ class WorkerInteractor:
             "runtest_protocol_complete", item_index=self.item_index, duration=duration
         )
 
+    def pytest_collection_modifyitems(self, session, config, items):
+        # add the group name to nodeid as suffix if --dist=loadgroup
+        if config.getvalue("loadgroup"):
+            for item in items:
+                mark = item.get_closest_marker("xdist_group")
+                if not mark:
+                    continue
+                gname = (
+                    mark.args[0]
+                    if len(mark.args) > 0
+                    else mark.kwargs.get("name", "default")
+                )
+                item._nodeid = "{}@{}".format(item.nodeid, gname)
+
     @pytest.hookimpl
     def pytest_collection_finish(self, session):
         try:
@@ -236,6 +250,7 @@ def remote_initconfig(option_dict, args):
 
 
 def setup_config(config, basetemp):
+    config.option.loadgroup = config.getvalue("dist") == "loadgroup"
     config.option.looponfail = False
     config.option.usepdb = False
     config.option.dist = "no"
diff --git a/src/xdist/scheduler/__init__.py b/src/xdist/scheduler/__init__.py
index 06ba6b7..ab2e830 100644
--- a/src/xdist/scheduler/__init__.py
+++ b/src/xdist/scheduler/__init__.py
@@ -2,3 +2,4 @@ from xdist.scheduler.each import EachScheduling  # noqa
 from xdist.scheduler.load import LoadScheduling  # noqa
 from xdist.scheduler.loadfile import LoadFileScheduling  # noqa
 from xdist.scheduler.loadscope import LoadScopeScheduling  # noqa
+from xdist.scheduler.loadgroup import LoadGroupScheduling  # noqa
diff --git a/src/xdist/scheduler/loadgroup.py b/src/xdist/scheduler/loadgroup.py
new file mode 100644
index 0000000..072f64a
--- /dev/null
+++ b/src/xdist/scheduler/loadgroup.py
@@ -0,0 +1,54 @@
+from .loadscope import LoadScopeScheduling
+from py.log import Producer
+
+
+class LoadGroupScheduling(LoadScopeScheduling):
+    """Implement load scheduling across nodes, but grouping test by xdist_group mark.
+
+    This class behaves very much like LoadScopeScheduling, but it groups tests by xdist_group mark
+    instead of the module or class to which they belong to.
+    """
+
+    def __init__(self, config, log=None):
+        super().__init__(config, log)
+        if log is None:
+            self.log = Producer("loadgroupsched")
+        else:
+            self.log = log.loadgroupsched
+
+    def _split_scope(self, nodeid):
+        """Determine the scope (grouping) of a nodeid.
+
+        There are usually 3 cases for a nodeid::
+
+            example/loadsuite/test/test_beta.py::test_beta0
+            example/loadsuite/test/test_delta.py::Delta1::test_delta0
+            example/loadsuite/epsilon/__init__.py::epsilon.epsilon
+
+        #. Function in a test module.
+        #. Method of a class in a test module.
+        #. Doctest in a function in a package.
+
+        With loadgroup, two cases are added::
+
+            example/loadsuite/test/test_beta.py::test_beta0
+            example/loadsuite/test/test_delta.py::Delta1::test_delta0
+            example/loadsuite/epsilon/__init__.py::epsilon.epsilon
+            example/loadsuite/test/test_gamma.py::test_beta0@gname
+            example/loadsuite/test/test_delta.py::Gamma1::test_gamma0@gname
+
+        This function will group tests with the scope determined by splitting the first ``@``
+        from the right. That is, test will be grouped in a single work unit when they have
+        same group name. In the above example, scopes will be::
+
+            example/loadsuite/test/test_beta.py::test_beta0
+            example/loadsuite/test/test_delta.py::Delta1::test_delta0
+            example/loadsuite/epsilon/__init__.py::epsilon.epsilon
+            gname
+            gname
+        """
+        if nodeid.rfind("@") > nodeid.rfind("]"):
+            # check the index of ']' to avoid the case: parametrize mark value has '@'
+            return nodeid.split("@")[-1]
+        else:
+            return nodeid


Test patch:
diff --git a/testing/acceptance_test.py b/testing/acceptance_test.py
index 31d634b..c139197 100644
--- a/testing/acceptance_test.py
+++ b/testing/acceptance_test.py
@@ -1326,6 +1326,134 @@ class TestFileScope:
         assert c1 == c2
 
 
+class TestGroupScope:
+    def test_by_module(self, testdir):
+        test_file = """
+            import pytest
+            class TestA:
+                @pytest.mark.xdist_group(name="xdist_group")
+                @pytest.mark.parametrize('i', range(5))
+                def test(self, i):
+                    pass
+        """
+        testdir.makepyfile(test_a=test_file, test_b=test_file)
+        result = testdir.runpytest("-n2", "--dist=loadgroup", "-v")
+        test_a_workers_and_test_count = get_workers_and_test_count_by_prefix(
+            "test_a.py::TestA", result.outlines
+        )
+        test_b_workers_and_test_count = get_workers_and_test_count_by_prefix(
+            "test_b.py::TestA", result.outlines
+        )
+
+        assert (
+            test_a_workers_and_test_count
+            in (
+                {"gw0": 5},
+                {"gw1": 0},
+            )
+            or test_a_workers_and_test_count in ({"gw0": 0}, {"gw1": 5})
+        )
+        assert (
+            test_b_workers_and_test_count
+            in (
+                {"gw0": 5},
+                {"gw1": 0},
+            )
+            or test_b_workers_and_test_count in ({"gw0": 0}, {"gw1": 5})
+        )
+        assert (
+            test_a_workers_and_test_count.items()
+            == test_b_workers_and_test_count.items()
+        )
+
+    def test_by_class(self, testdir):
+        testdir.makepyfile(
+            test_a="""
+            import pytest
+            class TestA:
+                @pytest.mark.xdist_group(name="xdist_group")
+                @pytest.mark.parametrize('i', range(10))
+                def test(self, i):
+                    pass
+            class TestB:
+                @pytest.mark.xdist_group(name="xdist_group")
+                @pytest.mark.parametrize('i', range(10))
+                def test(self, i):
+                    pass
+        """
+        )
+        result = testdir.runpytest("-n2", "--dist=loadgroup", "-v")
+        test_a_workers_and_test_count = get_workers_and_test_count_by_prefix(
+            "test_a.py::TestA", result.outlines
+        )
+        test_b_workers_and_test_count = get_workers_and_test_count_by_prefix(
+            "test_a.py::TestB", result.outlines
+        )
+
+        assert (
+            test_a_workers_and_test_count
+            in (
+                {"gw0": 10},
+                {"gw1": 0},
+            )
+            or test_a_workers_and_test_count in ({"gw0": 0}, {"gw1": 10})
+        )
+        assert (
+            test_b_workers_and_test_count
+            in (
+                {"gw0": 10},
+                {"gw1": 0},
+            )
+            or test_b_workers_and_test_count in ({"gw0": 0}, {"gw1": 10})
+        )
+        assert (
+            test_a_workers_and_test_count.items()
+            == test_b_workers_and_test_count.items()
+        )
+
+    def test_module_single_start(self, testdir):
+        test_file1 = """
+            import pytest
+            @pytest.mark.xdist_group(name="xdist_group")
+            def test():
+                pass
+        """
+        test_file2 = """
+            import pytest
+            def test_1():
+                pass
+            @pytest.mark.xdist_group(name="xdist_group")
+            def test_2():
+                pass
+        """
+        testdir.makepyfile(test_a=test_file1, test_b=test_file1, test_c=test_file2)
+        result = testdir.runpytest("-n2", "--dist=loadgroup", "-v")
+        a = get_workers_and_test_count_by_prefix("test_a.py::test", result.outlines)
+        b = get_workers_and_test_count_by_prefix("test_b.py::test", result.outlines)
+        c = get_workers_and_test_count_by_prefix("test_c.py::test_2", result.outlines)
+
+        assert a.keys() == b.keys() and b.keys() == c.keys()
+
+    def test_with_two_group_names(self, testdir):
+        test_file = """
+            import pytest
+            @pytest.mark.xdist_group(name="group1")
+            def test_1():
+                pass
+            @pytest.mark.xdist_group("group2")
+            def test_2():
+                pass
+        """
+        testdir.makepyfile(test_a=test_file, test_b=test_file)
+        result = testdir.runpytest("-n2", "--dist=loadgroup", "-v")
+        a_1 = get_workers_and_test_count_by_prefix("test_a.py::test_1", result.outlines)
+        a_2 = get_workers_and_test_count_by_prefix("test_a.py::test_2", result.outlines)
+        b_1 = get_workers_and_test_count_by_prefix("test_b.py::test_1", result.outlines)
+        b_2 = get_workers_and_test_count_by_prefix("test_b.py::test_2", result.outlines)
+
+        assert a_1.keys() == b_1.keys() and a_2.keys() == b_2.keys()
+
+
 class TestLocking:
     _test_content = """
     class TestClassName%s(object):