From 7d099b65c246844036575a1d3802d0d244bf3ea3 Mon Sep 17 00:00:00 2001
From: Michael Elovskikh <wronglink@yandex-team.ru>
Date: Sun, 2 Oct 2016 15:46:34 +0500
Subject: [PATCH] Allow custom scheduler class

---
 xdist/dsession.py | 24 +++++++++++++++---------
 xdist/plugin.py   |  4 ++++
 xdist/remote.py   |  1 +
 3 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/xdist/dsession.py b/xdist/dsession.py
index 46f76a9..8e87fa4 100644
--- a/xdist/dsession.py
+++ b/xdist/dsession.py
@@ -1,4 +1,5 @@
 import difflib
+import importlib
 import itertools
 from _pytest.runner import CollectReport
 
@@ -24,7 +25,7 @@ class EachScheduling:
     assigned the remaining items from the removed node.
     """
 
-    def __init__(self, numnodes, log=None):
+    def __init__(self, numnodes, log=None, config=None):
         self.numnodes = numnodes
         self.node2collection = {}
         self.node2pending = {}
@@ -485,6 +486,17 @@ class DSession:
             self.trdist = TerminalDistReporter(config)
             config.pluginmanager.register(self.trdist, "terminaldistreporter")
 
+        self.sched = None
+        dist = self.config.getvalue("dist")
+        if dist == "load":
+            self.ScheduleCls = LoadScheduling
+        elif dist == "each":
+            self.ScheduleCls = EachScheduling
+        else:
+            mod_name, cls_name = dist.rsplit('.', 1)
+            mod = importlib.import_module(mod_name)
+            self.ScheduleCls = getattr(mod, cls_name)
+
     @property
     def session_finished(self):
         """Return True if the distributed session has finished
@@ -523,15 +535,9 @@ class DSession:
 
     def pytest_runtestloop(self):
         numnodes = len(self.nodemanager.specs)
-        dist = self.config.getvalue("dist")
-        if dist == "load":
-            self.sched = LoadScheduling(numnodes, log=self.log,
-                                        config=self.config)
-        elif dist == "each":
-            self.sched = EachScheduling(numnodes, log=self.log)
-        else:
-            assert 0, dist
         self.shouldstop = False
+        self.sched = self.ScheduleCls(numnodes, log=self.log,
+                                      config=self.config)
         while not self.session_finished:
             self.loop_once()
             if self.shouldstop:
diff --git a/xdist/plugin.py b/xdist/plugin.py
index dc721da..77c81e8 100644
--- a/xdist/plugin.py
+++ b/xdist/plugin.py
@@ -37,6 +37,8 @@ def pytest_addoption(parser):
               "each: send each test to each available environment.\n\n"
               "load: send each test to available environment.\n\n"
               "(default) no: run tests inprocess, don't distribute."))
+    group._addoption('--dc', dest='distcustom',
+                    help="Custom scheduler implementation")
     group._addoption(
         '--tx', dest="tx", action="append", default=[],
         metavar="xspec",
@@ -99,6 +101,8 @@ def pytest_cmdline_main(config):
         config.option.tx = ['popen'] * config.option.numprocesses
     if config.option.distload:
         config.option.dist = "load"
+    if config.option.distcustom:
+        config.option.dist = config.option.distcustom
     val = config.getvalue
     if not val("collectonly"):
         usepdb = config.getoption('usepdb')  # a core option
diff --git a/xdist/remote.py b/xdist/remote.py
index 40bbeac..31dcc19 100644
--- a/xdist/remote.py
+++ b/xdist/remote.py
@@ -131,6 +131,7 @@ def remote_initconfig(option_dict, args):
     config.option.usepdb = False
     config.option.dist = "no"
     config.option.distload = False
+    config.option.distcustom = None
     config.option.numprocesses = None
     config.args = args
     return config