loosen legacy scheduler check

elastic · Nov 5, 2020 · 9b72a3c · 9b72a3c
1 parent 86c8025
commit 9b72a3c
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/esrally/driver/scheduler.py b/esrally/driver/scheduler.py
@@ -114,7 +114,7 @@ def is_legacy_scheduler(scheduler_class):
     target throughput.
     """
     constructor_params = inspect.signature(scheduler_class.__init__).parameters
-    return len(constructor_params) == 2 and "params" in constructor_params
+    return len(constructor_params) >= 2 and "params" in constructor_params
 
 
 def is_simple_scheduler(scheduler_class):

diff --git a/tests/driver/scheduler_test.py b/tests/driver/scheduler_test.py
@@ -126,8 +126,14 @@ class LegacyScheduler:
         def __init__(self, params):
             pass
 
+    class LegacySchedulerWithAdditionalArgs:
+        # pylint: disable=unused-variable
+        def __init__(self, params, my_default_param=True):
+            pass
+
     def test_detects_legacy_scheduler(self):
         self.assertTrue(scheduler.is_legacy_scheduler(SchedulerCategorizationTests.LegacyScheduler))
+        self.assertTrue(scheduler.is_legacy_scheduler(SchedulerCategorizationTests.LegacySchedulerWithAdditionalArgs))
 
     def test_a_regular_scheduler_is_not_a_legacy_scheduler(self):
         self.assertFalse(scheduler.is_legacy_scheduler(scheduler.DeterministicScheduler))