From d59db061495ee06e84e9ed513dad8364cec01ec5 Mon Sep 17 00:00:00 2001
From: Toan Vuong <toan.vuong@hyperscience.com>
Date: Wed, 12 Jun 2024 16:32:22 -0700
Subject: [PATCH] Retain input ordering in `loadscope`

* Optionally retain input ordering in loadscope for tests where relative
  ordering matters. i.e. guarantee that, given [input 1, input 2], input
  2 never runs before input 1. On any given worker, either input 1 has
  ran before input 2, or input 1 has never and will never run on this
  worker.
  Closes #1083.
---
 src/xdist/plugin.py              | 11 +++++++++++
 src/xdist/scheduler/loadscope.py | 14 +++++++++-----
 testing/acceptance_test.py       | 16 ++++++++++++++++
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/src/xdist/plugin.py b/src/xdist/plugin.py
index f670d9de..e5b7b901 100644
--- a/src/xdist/plugin.py
+++ b/src/xdist/plugin.py
@@ -127,6 +127,17 @@ def pytest_addoption(parser: pytest.Parser) -> None:
             "(default) no: Run tests inprocess, don't distribute."
         ),
     )
+    group.addoption(
+        "--no-loadscope-reorder",
+        action="store_true",
+        dest="noloadscopenoreorder",
+        default=False,
+        help=(
+            "Reorders tests when used in conjunction with loadscope.\n"
+            "Will order tests by number of tests per scope as a best-effort"
+            " attempt to evenly distribute scopes across all workers."
+        )
+    )
     group.addoption(
         "--tx",
         dest="tx",
diff --git a/src/xdist/scheduler/loadscope.py b/src/xdist/scheduler/loadscope.py
index a4d63b29..bb2e5b1a 100644
--- a/src/xdist/scheduler/loadscope.py
+++ b/src/xdist/scheduler/loadscope.py
@@ -371,11 +371,15 @@ def schedule(self) -> None:
             work_unit = unsorted_workqueue.setdefault(scope, {})
             work_unit[nodeid] = False
 
-        # Insert tests scopes into work queue ordered by number of tests.
-        for scope, nodeids in sorted(
-            unsorted_workqueue.items(), key=lambda item: -len(item[1])
-        ):
-            self.workqueue[scope] = nodeids
+        if self.config.option.noloadscopenoreorder:
+            for scope, nodeids in unsorted_workqueue.items():
+                self.workqueue[scope] = nodeids
+        else:
+            # Insert tests scopes into work queue ordered by number of tests.
+            for scope, nodeids in sorted(
+                unsorted_workqueue.items(), key=lambda item: -len(item[1])
+            ):
+                self.workqueue[scope] = nodeids
 
         # Avoid having more workers than work
         extra_nodes = len(self.nodes) - len(self.workqueue)
diff --git a/testing/acceptance_test.py b/testing/acceptance_test.py
index 3ef10cc9..32146b6e 100644
--- a/testing/acceptance_test.py
+++ b/testing/acceptance_test.py
@@ -1254,6 +1254,22 @@ def test(i):
             "test_b.py::test", result.outlines
         ) == {"gw0": 20}
 
+    def test_workqueue_ordered_by_input(self, pytester: pytest.Pytester) -> None:
+        test_file = """
+            import pytest
+            @pytest.mark.parametrize('i', range({}))
+            def test(i):
+                pass
+        """
+        pytester.makepyfile(test_a=test_file.format(10), test_b=test_file.format(20))
+        result = pytester.runpytest("-n2", "--dist=loadscope", "--no-loadscope-reorder", "-v")
+        assert get_workers_and_test_count_by_prefix(
+            "test_a.py::test", result.outlines
+        ) == {"gw0": 10}
+        assert get_workers_and_test_count_by_prefix(
+            "test_b.py::test", result.outlines
+        ) == {"gw1": 20}
+
     def test_module_single_start(self, pytester: pytest.Pytester) -> None:
         """Fix test suite never finishing in case all workers start with a single test (#277)."""
         test_file1 = """