Create new dist option 'loadgroup'

Author: baekdohyeop

src/xdist/dsession.py

@@ -7,6 +7,7 @@
     LoadScheduling,
     LoadScopeScheduling,
     LoadFileScheduling,
+    LoadGroupScheduling,
 )
 
 
@@ -100,6 +101,7 @@ def pytest_xdist_make_scheduler(self, config, log):
             "load": LoadScheduling,
             "loadscope": LoadScopeScheduling,
             "loadfile": LoadFileScheduling,
+            "loadgroup": LoadGroupScheduling,
         }
         return schedulers[dist](config, log)
 

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,8 @@ 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: load balance by sending any pending test or test group"
+            " to any available enviroment.\n\n"
             "(default) no: run tests inprocess, don't distribute."
         ),
     )
@@ -204,6 +206,12 @@ def pytest_configure(config):
         config.issue_config_time_warning(warning, 2)
         config.option.forked = True
 
+    config_line = (
+        "xgroup: 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):

src/xdist/remote.py

@@ -116,6 +116,20 @@ def run_one_test(self, torun):
             "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:
+                try:
+                    mark = item.get_closest_marker("xgroup")
+                except AttributeError:
+                    mark = item.get_marker("xgroup")
+
+                if mark:
+                    gname = mark.kwargs.get("name")
+                    if gname:
+                        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 = True if config.getvalue("dist") == "loadgroup" else False
     config.option.looponfail = False
     config.option.usepdb = False
     config.option.dist = "no"

src/xdist/scheduler/__init__.py

@@ -2,3 +2,4 @@
 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

src/xdist/scheduler/loadgroup.py

@@ -0,0 +1,67 @@
+from .loadscope import LoadScopeScheduling
+from py.log import Producer
+
+
+class LoadGroupScheduling(LoadScopeScheduling):
+    """Implement load scheduling across nodes, but grouping test only has group mark.
+
+    This distributes the tests collected across all nodes so each test is run
+    just once.  All nodes collect and submit the list of tests and when all
+    collections are received it is verified they are identical collections.
+    Then the collection gets divided up in work units, grouped by group mark
+    (If there is no group mark, it is itself a group.), and those work units
+    et submitted to nodes. Whenever a node finishes an item, it calls
+    ``.mark_test_complete()`` which will trigger the scheduler to assign more
+    work units if the number of pending tests for the node falls below a low-watermark.
+
+    When created, ``numnodes`` defines how many nodes are expected to submit a
+    collection. This is used to know when all nodes have finished collection.
+
+    This class behaves very much like LoadScopeScheduling,
+    but with a itself or group(by marked) scope.
+    """
+
+    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

testing/acceptance_test.py

@@ -1326,6 +1326,115 @@ def test_2():
         assert c1 == c2
 
 
+class TestGroupScope:
+    def test_by_module(self, testdir):
+        test_file = """
+            import pytest
+            class TestA:
+                @pytest.mark.xgroup(name="xgroup")
+                @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.xgroup(name="xgroup")
+                @pytest.mark.parametrize('i', range(10))
+                def test(self, i):
+                    pass
+            class TestB:
+                @pytest.mark.xgroup(name="xgroup")
+                @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.xgroup(name="xgroup")
+            def test():
+                pass
+        """
+        test_file2 = """
+            import pytest
+            def test_1():
+                pass
+            @pytest.mark.xgroup(name="xgroup")
+            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()
+
+
 class TestLocking:
     _test_content = """
     class TestClassName%s(object):