SERVER-41802: Properly handle max_sub_suite when splitting tasks

(cherry picked from commit 0a0625f43e4fa34302c7bc4757b0f273782925b5)
(cherry picked from commit 6dad28809ae9d59f82f6673c18023c798653a2ff)

4 files changed, 457 insertions, 307 deletions

diff --git a/buildscripts/evergreen_generate_resmoke_tasks.py b/buildscripts/evergreen_generate_resmoke_tasks.py
index 11b731f963c..d9a15215927 100755
--- a/buildscripts/evergreen_generate_resmoke_tasks.py
+++ b/buildscripts/evergreen_generate_resmoke_tasks.py
@@ -8,20 +8,22 @@ to attempt to keep the task runtime under a specified amount.
 
 from __future__ import absolute_import
 
-import argparse
 import datetime
+from datetime import timedelta
 import logging
 import math
 import os
 import re
 import sys
 from collections import defaultdict
-from collections import namedtuple
 from distutils.util import strtobool  # pylint: disable=no-name-in-module
 
+import click
 import requests
+import structlog
 import yaml
 
+from evergreen.api import RetryingEvergreenApi
 from shrub.config import Configuration
 from shrub.command import CommandDefinition
 from shrub.operations import CmdTimeoutUpdate
@@ -33,53 +35,145 @@ from shrub.variant import TaskSpec
 if __name__ == "__main__" and __package__ is None:
     sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-import buildscripts.client.evergreen as evergreen  # pylint: disable=wrong-import-position
 import buildscripts.resmokelib.suitesconfig as suitesconfig  # pylint: disable=wrong-import-position
 import buildscripts.util.read_config as read_config  # pylint: disable=wrong-import-position
 import buildscripts.util.taskname as taskname  # pylint: disable=wrong-import-position
 import buildscripts.util.testname as testname  # pylint: disable=wrong-import-position
 
-LOGGER = logging.getLogger(__name__)
+LOGGER = structlog.getLogger(__name__)
 
 TEST_SUITE_DIR = os.path.join("buildscripts", "resmokeconfig", "suites")
 CONFIG_DIR = "generated_resmoke_config"
-MIN_TIMEOUT_SECONDS = 5 * 60
+CONFIG_FILE = "./.evergreen.yml"
+MIN_TIMEOUT_SECONDS = int(timedelta(minutes=5).total_seconds())
+LOOKBACK_DURATION_DAYS = 14
+GEN_SUFFIX = "_gen"
 
 HEADER_TEMPLATE = """# DO NOT EDIT THIS FILE. All manual edits will be lost.
 # This file was generated by {file} from
 # {suite_file}.
 """
 
-ConfigOptions = namedtuple("ConfigOptions", [
-    "build_id",
-    "depends_on",
+REQUIRED_CONFIG_KEYS = {
+    "build_variant",
     "fallback_num_sub_suites",
-    "is_patch",
-    "large_distro_name",
-    "max_sub_suites",
     "project",
-    "repeat_suites",
-    "resmoke_args",
-    "resmoke_jobs_max",
-    "run_multiple_jobs",
-    "suite",
-    "target_resmoke_time",
-    "task",
-    "use_default_timeouts",
-    "use_large_distro",
-    "use_multiversion",
-    "variant",
-])
-
-
-def enable_logging():
+    "task_id",
+    "task_name",
+}
+
+DEFAULT_CONFIG_VALUES = {
+    "resmoke_args": "",
+    "resmoke_repeat_suites": 1,
+    "run_multiple_jobs": "true",
+    "target_resmoke_time": 60,
+    "use_default_timeouts": False,
+    "use_large_distro": False,
+}
+
+CONFIG_FORMAT_FN = {
+    "fallback_num_sub_suites": int,
+    "max_sub_suites": int,
+    "target_resmoke_time": int,
+}
+
+
+class ConfigOptions(object):
+    """Retrieve configuration from a config file."""
+
+    def __init__(self, config, required_keys=None, defaults=None, formats=None):
+        """
+        Create an instance of ConfigOptions.
+
+        :param config: Dictionary of configuration to use.
+        :param required_keys: Set of keys required by this config.
+        :param defaults: Dict of default values for keys.
+        :param formats: Dict with functions to format values before returning.
+        """
+        self.config = config
+        self.required_keys = required_keys if required_keys else set()
+        self.default_values = defaults if defaults else {}
+        self.formats = formats if formats else {}
+
+    @classmethod
+    def from_file(cls, filepath, required_keys, defaults, formats):
+        """
+        Create an instance of ConfigOptions based on the given config file.
+
+        :param filepath: Path to file containing configuration.
+        :param required_keys: Set of keys required by this config.
+        :param defaults: Dict of default values for keys.
+        :param formats: Dict with functions to format values before returning.
+        :return: Instance of ConfigOptions.
+        """
+        return cls(read_config.read_config_file(filepath), required_keys, defaults, formats)
+
+    @property
+    def depends_on(self):
+        """List of dependencies specified."""
+        return split_if_exists(self._lookup(self.config, "depends_on"))
+
+    @property
+    def is_patch(self):
+        """Is this running in a patch build."""
+        patch = self.config.get("is_patch")
+        if patch:
+            return strtobool(patch)
+        return None
+
+    @property
+    def repeat_suites(self):
+        """How many times should the suite be repeated."""
+        return int(self.resmoke_repeat_suites)
+
+    @property
+    def suite(self):
+        """Return test suite is being run."""
+        return self.config.get("suite", self.task)
+
+    @property
+    def task(self):
+        """Return task being run."""
+        return remove_gen_suffix(self.task_name)
+
+    @property
+    def variant(self):
+        """Return build variant is being run on."""
+        return self.build_variant
+
+    def _lookup(self, config, item):
+        if item not in config:
+            if item in self.required_keys:
+                raise KeyError("{item} must be specified in configuration.".format(item=item))
+            return self.default_values.get(item, None)
+
+        if item in self.formats and item in config:
+            return self.formats[item](config[item])
+
+        return config.get(item, None)
+
+    def __getattr__(self, item):
+        """Determine the value of the given attribute."""
+        return self._lookup(self.config, item)
+
+    def __repr__(self):
+        """Provide a string representation of this object for debugging."""
+        required_values = [
+            "{key}: {value}".format(key=key, value=self.config[key]) for key in REQUIRED_CONFIG_KEYS
+        ]
+        return "ConfigOptions({options})".format(options=", ".join(required_values))
+
+
+def enable_logging(verbose):
     """Enable verbose logging for execution."""
 
+    level = logging.DEBUG if verbose else logging.INFO
     logging.basicConfig(
         format="[%(asctime)s - %(name)s - %(levelname)s] %(message)s",
-        level=logging.DEBUG,
+        level=level,
         stream=sys.stdout,
     )
+    structlog.configure(logger_factory=structlog.stdlib.LoggerFactory())
 
 
 def split_if_exists(str_to_split):
@@ -89,61 +183,22 @@ def split_if_exists(str_to_split):
     return None
 
 
-def get_config_options(cmd_line_options, config_file):
-    """
-    Get the configuration to use for generated tests.
+def remove_gen_suffix(task_name):
+    """Remove '_gen' suffix from task_name."""
+    if task_name.endswith(GEN_SUFFIX):
+        return task_name[:-4]
+    return task_name
 
-    Command line options override config file options.
 
-    :param cmd_line_options: Command line options specified.
-    :param config_file: config file to use.
-    :return: ConfigOptions to use.
+def string_contains_any_of_args(string, args):
     """
-    # pylint: disable=too-many-locals
-    config_file_data = read_config.read_config_file(config_file)
-
-    fallback_num_sub_suites = int(
-        read_config.get_config_value("fallback_num_sub_suites", cmd_line_options, config_file_data,
-                                     required=True))
-    max_sub_suites = read_config.get_config_value("max_sub_suites", cmd_line_options,
-                                                  config_file_data)
-    project = read_config.get_config_value("project", cmd_line_options, config_file_data,
-                                           required=True)
-    resmoke_args = read_config.get_config_value("resmoke_args", cmd_line_options, config_file_data,
-                                                default="")
-    resmoke_jobs_max = read_config.get_config_value("resmoke_jobs_max", cmd_line_options,
-                                                    config_file_data)
-    target_resmoke_time = int(
-        read_config.get_config_value("target_resmoke_time", cmd_line_options, config_file_data,
-                                     default=60))
-    run_multiple_jobs = read_config.get_config_value("run_multiple_jobs", cmd_line_options,
-                                                     config_file_data, default="true")
-    task = read_config.get_config_value("task", cmd_line_options, config_file_data, required=True)
-    suite = read_config.get_config_value("suite", cmd_line_options, config_file_data, default=task)
-    variant = read_config.get_config_value("build_variant", cmd_line_options, config_file_data,
-                                           required=True)
-    use_default_timeouts = read_config.get_config_value("use_default_timeouts", cmd_line_options,
-                                                        config_file_data, default=False)
-    use_large_distro = read_config.get_config_value("use_large_distro", cmd_line_options,
-                                                    config_file_data, default=False)
-    large_distro_name = read_config.get_config_value("large_distro_name", cmd_line_options,
-                                                     config_file_data)
-    use_multiversion = read_config.get_config_value("use_multiversion", cmd_line_options,
-                                                    config_file_data)
-    is_patch = read_config.get_config_value("is_patch", cmd_line_options, config_file_data)
-    if is_patch:
-        is_patch = strtobool(is_patch)
-    depends_on = split_if_exists(
-        read_config.get_config_value("depends_on", cmd_line_options, config_file_data))
-    build_id = read_config.get_config_value("build_id", cmd_line_options, config_file_data)
-    repeat_suites = int(
-        read_config.get_config_value("resmoke_repeat_suites", cmd_line_options, config_file_data,
-                                     default=1))
-
-    return ConfigOptions(build_id, depends_on, fallback_num_sub_suites, is_patch, large_distro_name,
-                         max_sub_suites, project, repeat_suites, resmoke_args, resmoke_jobs_max,
-                         run_multiple_jobs, suite, target_resmoke_time, task, use_default_timeouts,
-                         use_large_distro, use_multiversion, variant)
+    Return whether array contains any of a group of args.
+
+    :param string: String being checked.
+    :param args: Args being analyzed.
+    :return: True if any args are found in the string.
+    """
+    return any(arg in string for arg in args)
 
 
 def divide_remaining_tests_among_suites(remaining_tests_runtimes, suites):
@@ -166,12 +221,12 @@ def divide_tests_into_suites(tests_runtimes, max_time_seconds, max_suites=None):
     suites = []
     current_suite = Suite()
     last_test_processed = len(tests_runtimes)
-    LOGGER.debug("Determines suites for runtime: %ds", max_time_seconds)
+    LOGGER.debug("Determines suites for runtime", max_runtime_seconds=max_time_seconds)
     for idx, (test_file, runtime) in enumerate(tests_runtimes):
-        LOGGER.debug("Adding test %s, runtime %d", test_file, runtime)
+        LOGGER.debug("Adding test", test=test_file, test_runtime=runtime)
         if current_suite.get_runtime() + runtime > max_time_seconds:
-            LOGGER.debug("Runtime(%d) + new test(%d) > max(%d)", current_suite.get_runtime(),
-                         runtime, max_time_seconds)
+            LOGGER.debug("Runtime greater than boundary", suite_runtime=current_suite.get_runtime(),
+                         test_runtime=runtime, max_time=max_time_seconds)
             if current_suite.get_test_count() > 0:
                 suites.append(current_suite)
                 current_suite = Suite()
@@ -228,7 +283,7 @@ def generate_subsuite_file(source_suite_name, target_suite_name, roots=None, exc
     """
     source_file = os.path.join(TEST_SUITE_DIR, source_suite_name + ".yml")
     with open(source_file, "r") as fstream:
-        suite_config = yaml.load(fstream)
+        suite_config = yaml.safe_load(fstream)
 
     with open(os.path.join(CONFIG_DIR, target_suite_name + ".yml"), "w") as out:
         out.write(HEADER_TEMPLATE.format(file=__file__, suite_file=source_file))
@@ -324,13 +379,13 @@ class EvergreenConfigGenerator(object):
             cmd_timeout = CmdTimeoutUpdate()
             if max_test_runtime:
                 timeout = calculate_timeout(max_test_runtime, 3) * repeat_factor
-                LOGGER.debug("Setting timeout to: %d (max=%d, repeat=%d)", timeout,
-                             max_test_runtime, repeat_factor)
+                LOGGER.debug("Setting timeout", timeout=timeout, max_runtime=max_test_runtime,
+                             factor=repeat_factor)
                 cmd_timeout.timeout(timeout)
             if expected_suite_runtime:
                 exec_timeout = calculate_timeout(expected_suite_runtime, 3) * repeat_factor
-                LOGGER.debug("Setting exec_timeout to: %d (runtime=%d, repeat=%d)", exec_timeout,
-                             expected_suite_runtime, repeat_factor)
+                LOGGER.debug("Setting exec_timeout", exec_timeout=exec_timeout,
+                             suite_runtime=expected_suite_runtime, factor=repeat_factor)
                 cmd_timeout.exec_timeout(exec_timeout)
             commands.append(cmd_timeout.validate().resolve())
 
@@ -340,8 +395,8 @@ class EvergreenConfigGenerator(object):
 
     def _get_tasks_for_depends_on(self, dependent_task):
         return [
-            str(task["display_name"]) for task in self.build_tasks
-            if self._is_task_dependency(dependent_task, str(task["display_name"]))
+            str(task.display_name) for task in self.build_tasks
+            if self._is_task_dependency(dependent_task, str(task.display_name))
         ]
 
     def _add_dependencies(self, task):
@@ -359,7 +414,7 @@ class EvergreenConfigGenerator(object):
     def _generate_task(self, sub_suite_name, sub_task_name, max_test_runtime=None,
                        expected_suite_runtime=None):
         """Generate evergreen config for a resmoke task."""
-        LOGGER.debug("Generating task for: %s", sub_suite_name)
+        LOGGER.debug("Generating task", sub_suite=sub_suite_name)
         spec = TaskSpec(sub_task_name)
         self._set_task_distro(spec)
         self.task_specs.append(spec)
@@ -411,7 +466,7 @@ class EvergreenConfigGenerator(object):
 
     def generate_config(self):
         """Generate evergreen configuration."""
-        self.build_tasks = self.evg_api.tasks_by_build_id(self.options.build_id)
+        self.build_tasks = self.evg_api.tasks_by_build(self.options.build_id)
         self._generate_variant()
         return self.evg_config
 
@@ -437,9 +492,9 @@ class TestStats(object):
 
     def _add_stats(self, test_stats):
         """Add the statistics found in a document returned by the Evergreen test_stats/ endpoint."""
-        test_file = testname.normalize_test_file(test_stats["test_file"])
-        duration = test_stats["avg_duration_pass"]
-        num_run = test_stats["num_pass"]
+        test_file = testname.normalize_test_file(test_stats.test_file)
+        duration = test_stats.avg_duration_pass
+        num_run = test_stats.num_pass
         is_hook = testname.is_resmoke_hook(test_file)
         if is_hook:
             self._add_test_hook_stats(test_file, duration, num_run)
@@ -470,7 +525,10 @@ class TestStats(object):
     @staticmethod
     def _average(value_a, num_a, value_b, num_b):
         """Compute a weighted average of 2 values with associated numbers."""
-        return float(value_a * num_a + value_b * num_b) / (num_a + num_b)
+        divisor = num_a + num_b
+        if divisor == 0:
+            return 0
+        return float(value_a * num_a + value_b * num_b) / divisor
 
     def get_tests_runtimes(self):
         """Return the list of (test_file, runtime_in_secs) tuples ordered by decreasing runtime."""
@@ -525,64 +583,15 @@ class Suite(object):
         return len(self.tests)
 
 
-class Main(object):
+class GenerateSubSuites(object):
     """Orchestrate the execution of generate_resmoke_suites."""
 
-    def __init__(self, evergreen_api):
+    def __init__(self, evergreen_api, config_options):
         """Initialize the object."""
         self.evergreen_api = evergreen_api
-        self.options = None
-        self.config_options = None
+        self.config_options = config_options
         self.test_list = []
 
-    def parse_commandline(self):
-        """Parse the command line options and return the parsed data."""
-        parser = argparse.ArgumentParser(description=self.main.__doc__)
-
-        parser.add_argument("--expansion-file", dest="expansion_file", type=str,
-                            help="Location of expansions file generated by evergreen.")
-        parser.add_argument("--analysis-duration", dest="duration_days", default=14,
-                            help="Number of days to analyze.", type=int)
-        parser.add_argument("--execution-time", dest="target_resmoke_time", type=int,
-                            help="Target execution time (in minutes).")
-        parser.add_argument("--max-sub-suites", dest="max_sub_suites", type=int,
-                            help="Max number of suites to divide into.")
-        parser.add_argument("--fallback-num-sub-suites", dest="fallback_num_sub_suites", type=int,
-                            help="The number of suites to divide into if the Evergreen test "
-                            "statistics are not available.")
-        parser.add_argument("--project", dest="project", help="The Evergreen project to analyse.")
-        parser.add_argument("--resmoke-args", dest="resmoke_args",
-                            help="Arguments to pass to resmoke calls.")
-        parser.add_argument("--resmoke-jobs_max", dest="resmoke_jobs_max",
-                            help="Number of resmoke jobs to invoke.")
-        parser.add_argument("--suite", dest="suite",
-                            help="Name of suite being split(defaults to task_name)")
-        parser.add_argument("--run-multiple-jobs", dest="run_multiple_jobs",
-                            help="Should resmoke run multiple jobs")
-        parser.add_argument("--task-name", dest="task", help="Name of task to split.")
-        parser.add_argument("--variant", dest="build_variant",
-                            help="Build variant being run against.")
-        parser.add_argument("--use-large-distro", dest="use_large_distro",
-                            help="Should subtasks use large distros.")
-        parser.add_argument("--large-distro-name", dest="large_distro_name",
-                            help="Name of large distro.")
-        parser.add_argument("--use-default_timeouts", dest="use_default_timeouts",
-                            help="When specified, do not override timeouts based on test history.")
-        parser.add_argument("--use-multiversion", dest="use_multiversion",
-                            help="Task path suffix for multiversion generated tasks.")
-        parser.add_argument("--is-patch", dest="is_patch", help="Is this part of a patch build.")
-        parser.add_argument("--depends-on", dest="depends_on",
-                            help="Generate depends on for these tasks.")
-        parser.add_argument("--repeat-suites", dest="resmoke_repeat_suites",
-                            help="Repeat each suite the specified number of times.")
-        parser.add_argument("--verbose", dest="verbose", action="store_true", default=False,
-                            help="Enable verbose logging.")
-
-        options = parser.parse_args()
-
-        self.config_options = get_config_options(options, options.expansion_file)
-        return options
-
     def calculate_suites(self, start_date, end_date):
         """Divide tests into suites based on statistics for the provided period."""
         try:
@@ -609,10 +618,10 @@ class Main(object):
         # pylint: disable=too-many-arguments
 
         days = (end_date - start_date).days
-        return self.evergreen_api.test_stats(project, after_date=start_date.strftime("%Y-%m-%d"),
-                                             before_date=end_date.strftime("%Y-%m-%d"),
-                                             tasks=[task], variants=[variant], group_by="test",
-                                             group_num_days=days)
+        return self.evergreen_api.test_stats_by_project(
+            project, after_date=start_date.strftime("%Y-%m-%d"),
+            before_date=end_date.strftime("%Y-%m-%d"), tasks=[task], variants=[variant],
+            group_by="test", group_num_days=days)
 
     def calculate_suites_from_evg_stats(self, data, execution_time_secs):
         """Divide tests into suites that can be run in less than the specified execution time."""
@@ -622,7 +631,7 @@ class Main(object):
             return self.calculate_fallback_suites()
         self.test_list = [info[0] for info in tests_runtimes]
         return divide_tests_into_suites(tests_runtimes, execution_time_secs,
-                                        self.options.max_sub_suites)
+                                        self.config_options.max_sub_suites)
 
     def filter_existing_tests(self, tests_runtimes):
         """Filter out tests that do not exist in the filesystem."""
@@ -650,25 +659,18 @@ class Main(object):
         with open(os.path.join(CONFIG_DIR, task + ".json"), "w") as file_handle:
             file_handle.write(evg_config.to_json())
 
-    def main(self):
+    def run(self):
         """Generate resmoke suites that run within a specified target execution time."""
+        LOGGER.debug("config options", config_options=self.config_options)
 
-        options = self.parse_commandline()
-        self.options = options
-
-        if options.verbose:
-            enable_logging()
-
-        LOGGER.debug("Starting execution for options %s", options)
-        LOGGER.debug("config options %s", self.config_options)
         end_date = datetime.datetime.utcnow().replace(microsecond=0)
-        start_date = end_date - datetime.timedelta(days=options.duration_days)
+        start_date = end_date - datetime.timedelta(days=LOOKBACK_DURATION_DAYS)
 
         prepare_directory_for_suite(CONFIG_DIR)
 
         suites = self.calculate_suites(start_date, end_date)
 
-        LOGGER.debug("Creating %d suites for %s", len(suites), self.config_options.task)
+        LOGGER.debug("Creating suites", num_suites=len(suites), task=self.config_options.task)
 
         render_suite(suites, self.config_options.suite)
         render_misc_suite(self.test_list, self.config_options.suite)
@@ -676,5 +678,29 @@ class Main(object):
         self.write_evergreen_configuration(suites, self.config_options.task)
 
 
+@click.command()
+@click.option("--expansion-file", type=str, required=True,
+              help="Location of expansions file generated by evergreen.")
+@click.option("--evergreen-config", type=str, default=CONFIG_FILE,
+              help="Location of evergreen configuration file.")
+@click.option("--verbose", is_flag=True, default=False, help="Enable verbose logging.")
+def main(expansion_file, evergreen_config, verbose):
+    """
+    Create a configuration for generate tasks to create sub suites for the specified resmoke suite.
+
+    The `--expansion-file` should contain all the configuration needed to generate the tasks.
+    \f
+    :param expansion_file: Configuration file.
+    :param evergreen_config: Evergreen configuration file.
+    :param verbose: Use verbose logging.
+    """
+    enable_logging(verbose)
+    evg_api = RetryingEvergreenApi.get_api(config_file=evergreen_config)
+    config_options = ConfigOptions.from_file(expansion_file, REQUIRED_CONFIG_KEYS,
+                                             DEFAULT_CONFIG_VALUES, CONFIG_FORMAT_FN)
+
+    GenerateSubSuites(evg_api, config_options).run()
+
+
 if __name__ == "__main__":
-    Main(evergreen.get_evergreen_apiv2()).main()
+    main()  # pylint: disable=no-value-for-parameter
diff --git a/buildscripts/requirements.txt b/buildscripts/requirements.txt
index dc129a6991c..d83d273f6fa 100644
--- a/buildscripts/requirements.txt
+++ b/buildscripts/requirements.txt
@@ -3,7 +3,7 @@ cryptography >= 2.0.0
 jira == 1.0.10
 pyjwt == 1.5.3
 # Other
-pyyaml == 3.11
+pyyaml == 5.1
 unittest-xml-reporting == 2.1.0
 # Linters
 yapf == 0.21.0
@@ -18,5 +18,8 @@ pydocstyle == 2.1.1
 cheetah3 == 3.0.0; python_version < "3"
 jinja2 == 2.10
 mock == 2.0.0
+click ~= 7.0
+evergreen.py == 0.3.9
 shrub.py == 0.2.0
+structlog ~= 19.1.0
 psutil == 5.4.3
diff --git a/buildscripts/resmokelib/requirements.txt b/buildscripts/resmokelib/requirements.txt
index 9d7bd449144..beb7432ecf9 100644
--- a/buildscripts/resmokelib/requirements.txt
+++ b/buildscripts/resmokelib/requirements.txt
@@ -2,6 +2,6 @@ mock == 2.0.0 ; python_version < "3"
 pymongo >= 3.0
 pypiwin32 == 219 ; sys_platform == "win32" and python_version < "3"
 pypiwin32 == 223 ; sys_platform == "win32" and python_version > "3"
-PyYAML == 3.11
+PyYAML == 5.1
 requests >= 2.16.1
 subprocess32 >= 3.2.7 ; os_name == "posix" and python_version < "3"
diff --git a/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py b/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
index 09757a85ad2..bd8f4b0a4a9 100644
--- a/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
+++ b/buildscripts/tests/test_evergreen_generate_resmoke_tasks.py
@@ -10,11 +10,9 @@ import unittest
 import requests
 import yaml
 
-from mock import patch, mock_open, call, Mock
+from mock import patch, mock_open, call, MagicMock
 
-from buildscripts import evergreen_generate_resmoke_tasks as grt
-from buildscripts.evergreen_generate_resmoke_tasks import render_suite, render_misc_suite, \
-    prepare_directory_for_suite
+from buildscripts import evergreen_generate_resmoke_tasks as under_test
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,no-self-use,protected-access
 
@@ -35,7 +33,7 @@ class TestTestStats(unittest.TestCase):
             self._make_evg_result("dir/test2.js", 1, 30),
             self._make_evg_result("dir/test1.js", 2, 25),
         ]
-        test_stats = grt.TestStats(evg_results)
+        test_stats = under_test.TestStats(evg_results)
         expected_runtimes = [
             ("dir/test2.js", 30),
             ("dir/test1.js", 20),
@@ -51,7 +49,7 @@ class TestTestStats(unittest.TestCase):
             self._make_evg_result("test3:CleanEveryN", 10, 30),
             self._make_evg_result("test3:CheckReplDBHash", 10, 35),
         ]
-        test_stats = grt.TestStats(evg_results)
+        test_stats = under_test.TestStats(evg_results)
         expected_runtimes = [
             ("dir/test3.js", 75),
             ("dir/test2.js", 30),
@@ -59,18 +57,100 @@ class TestTestStats(unittest.TestCase):
         ]
         self.assertEqual(expected_runtimes, test_stats.get_tests_runtimes())
 
+    def test_zero_runs(self):
+        evg_results = [
+            self._make_evg_result("dir/test1.js", 0, 0),
+            self._make_evg_result("dir/test1.js", 0, 0),
+        ]
+        test_stats = under_test.TestStats(evg_results)
+        expected_runtimes = [
+            ("dir/test1.js", 0),
+        ]
+        self.assertEqual(expected_runtimes, test_stats.get_tests_runtimes())
+
     @staticmethod
     def _make_evg_result(test_file="dir/test1.js", num_pass=0, duration=0):
-        return {
-            "test_file": test_file,
-            "task_name": "task1",
-            "variant": "variant1",
-            "distro": "distro1",
-            "date": _DATE,
-            "num_pass": num_pass,
-            "num_fail": 0,
-            "avg_duration_pass": duration,
-        }
+        mock = MagicMock()
+        mock.test_file = test_file
+        mock.task_name = "task1"
+        mock.variant = "variant1"
+        mock.distro = "distro1"
+        mock.date = _DATE
+        mock.num_pass = num_pass
+        mock.num_fail = 0
+        mock.avg_duration_pass = duration
+        return mock
+
+
+class TestConfigOptions(unittest.TestCase):
+    def test_lookup_missing_required_key_throws_exception(self):
+        config = {}
+        required_keys = {"key1"}
+
+        config_options = under_test.ConfigOptions(config, required_keys=required_keys)
+
+        with self.assertRaises(KeyError):
+            config_options.key1  # pylint: disable=pointless-statement
+
+    def test_default_options_use_default_value(self):
+        config = {}
+        defaults = {"key1": "value1"}
+
+        config_options = under_test.ConfigOptions(config, defaults=defaults)
+
+        self.assertEqual(defaults["key1"], config_options.key1)
+
+    def test_unknown_values_return_none(self):
+        config = {}
+
+        config_options = under_test.ConfigOptions(config)
+
+        self.assertIsNone(config_options.key1)
+
+    def test_set_values_are_used(self):
+        config = {"key1": "not_the_default_value"}
+        defaults = {"key1": "value1"}
+
+        config_options = under_test.ConfigOptions(config, defaults=defaults)
+
+        self.assertEqual(config["key1"], config_options.key1)
+
+    def test_depends_on_splits_values(self):
+        config = {"depends_on": "value1,value2,value3"}
+
+        config_options = under_test.ConfigOptions(config)
+
+        self.assertEqual(3, len(config_options.depends_on))
+
+    def test_suite_uses_task_value_if_no_suite(self):
+        config = {"task": "task_value"}
+
+        config_options = under_test.ConfigOptions(config)
+
+        self.assertEqual(config["task"], config_options.suite)
+
+    def test_suite_uses_suite_if_provided(self):
+        config = {"task": "task_value", "suite": "suite_value"}
+
+        config_options = under_test.ConfigOptions(config)
+
+        self.assertEqual(config["suite"], config_options.suite)
+
+    def test_task_uses_task_name(self):
+        config = {"task_name": "task_value"}
+
+        config_options = under_test.ConfigOptions(config)
+
+        self.assertEqual(config["task_name"], config_options.task)
+
+    def test_item_with_format_function_works(self):
+        config = {"number": "1"}
+        formats = {"number": int}
+
+        config_options = under_test.ConfigOptions(config, formats=formats)
+
+        self.assertEqual(1, config_options.number)
+        self.assertIsInstance(config_options.number, int)
 
 
 class DivideRemainingTestsAmongSuitesTest(unittest.TestCase):
@@ -85,19 +165,19 @@ class DivideRemainingTestsAmongSuitesTest(unittest.TestCase):
         return tests_runtimes
 
     def test_each_suite_gets_one_test(self):
-        suites = [grt.Suite(), grt.Suite(), grt.Suite()]
+        suites = [under_test.Suite(), under_test.Suite(), under_test.Suite()]
         tests_runtimes = self.generate_tests_runtimes(3)
 
-        grt.divide_remaining_tests_among_suites(tests_runtimes, suites)
+        under_test.divide_remaining_tests_among_suites(tests_runtimes, suites)
 
         for suite in suites:
             self.assertEqual(suite.get_test_count(), 1)
 
     def test_each_suite_gets_at_least_one_test(self):
-        suites = [grt.Suite(), grt.Suite(), grt.Suite()]
+        suites = [under_test.Suite(), under_test.Suite(), under_test.Suite()]
         tests_runtimes = self.generate_tests_runtimes(5)
 
-        grt.divide_remaining_tests_among_suites(tests_runtimes, suites)
+        under_test.divide_remaining_tests_among_suites(tests_runtimes, suites)
 
         total_tests = 0
         for suite in suites:
@@ -116,7 +196,7 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
             ("test3", 3),
         ]
 
-        suites = grt.divide_tests_into_suites(tests_runtimes, max_time)
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time)
         self.assertEqual(len(suites), 1)
         self.assertEqual(suites[0].get_test_count(), 3)
         self.assertEqual(suites[0].get_runtime(), 12)
@@ -129,7 +209,7 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
             ("test3", 3),
         ]
 
-        suites = grt.divide_tests_into_suites(tests_runtimes, max_time)
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time)
         self.assertEqual(len(suites), 3)
 
     def test_if_test_is_greater_than_max_it_goes_alone(self):
@@ -140,7 +220,7 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
             ("test3", 3),
         ]
 
-        suites = grt.divide_tests_into_suites(tests_runtimes, max_time)
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time)
         self.assertEqual(len(suites), 2)
         self.assertEqual(suites[0].get_test_count(), 1)
         self.assertEqual(suites[0].get_runtime(), 15)
@@ -156,7 +236,8 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
             ("test5", 3),
         ]
 
-        suites = grt.divide_tests_into_suites(tests_runtimes, max_time, max_suites=max_suites)
+        suites = under_test.divide_tests_into_suites(tests_runtimes, max_time,
+                                                     max_suites=max_suites)
         self.assertEqual(len(suites), max_suites)
         total_tests = 0
         for suite in suites:
@@ -166,7 +247,7 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
 
 class SuiteTest(unittest.TestCase):
     def test_adding_tests_increases_count_and_runtime(self):
-        suite = grt.Suite()
+        suite = under_test.Suite()
         suite.add_test("test1", 10)
         suite.add_test("test2", 12)
         suite.add_test("test3", 7)
@@ -176,7 +257,7 @@ class SuiteTest(unittest.TestCase):
         self.assertTrue(suite.should_overwrite_timeout())
 
     def test_suites_without_full_runtime_history_should_not_be_overridden(self):
-        suite = grt.Suite()
+        suite = under_test.Suite()
         suite.add_test("test1", 10)
         suite.add_test("test2", 0)
         suite.add_test("test3", 7)
@@ -186,7 +267,7 @@ class SuiteTest(unittest.TestCase):
 
 def create_suite(count=3, start=0):
     """ Create a suite with count tests."""
-    suite = grt.Suite()
+    suite = under_test.Suite()
     for i in range(start, start + count):
         suite.add_test("test{}".format(i), 1)
     return suite
@@ -196,25 +277,25 @@ class UpdateSuiteConfigTest(unittest.TestCase):
     def test_roots_are_updated(self):
         config = {"selector": {}}
 
-        updated_config = grt.update_suite_config(config, "root value")
+        updated_config = under_test.update_suite_config(config, "root value")
         self.assertEqual("root value", updated_config["selector"]["roots"])
 
     def test_excluded_files_not_included_if_not_specified(self):
         config = {"selector": {"excluded_files": "files to exclude"}}
 
-        updated_config = grt.update_suite_config(config, excludes=None)
+        updated_config = under_test.update_suite_config(config, excludes=None)
         self.assertNotIn("exclude_files", updated_config["selector"])
 
     def test_excluded_files_added_to_misc(self):
         config = {"selector": {}}
 
-        updated_config = grt.update_suite_config(config, excludes="files to exclude")
+        updated_config = under_test.update_suite_config(config, excludes="files to exclude")
         self.assertEqual("files to exclude", updated_config["selector"]["exclude_files"])
 
     def test_excluded_files_extended_in_misc(self):
         config = {"selector": {"exclude_files": ["file 0", "file 1"]}}
 
-        updated_config = grt.update_suite_config(config, excludes=["file 2", "file 3"])
+        updated_config = under_test.update_suite_config(config, excludes=["file 2", "file 3"])
         self.assertEqual(4, len(updated_config["selector"]["exclude_files"]))
         for exclude in ["file 0", "file 1", "file 2", "file 3"]:
             self.assertIn(exclude, updated_config["selector"]["exclude_files"])
@@ -239,18 +320,18 @@ class RenderSuites(unittest.TestCase):
 
         m = mock_open(read_data=yaml.dump({"selector": {"roots": [], "excludes": ["fixed"]}}))
         with patch(ns("open"), m, create=True):
-            render_suite(suites, "suite_name")
+            under_test.render_suite(suites, "suite_name")
         handle = m()
 
         # The other writes are for the headers.
         self.assertEquals(len(suites) * 2, handle.write.call_count)
         handle.write.assert_has_calls([call(e) for e in expected], any_order=True)
         calls = [
-            call(os.path.join(grt.TEST_SUITE_DIR, "suite_name.yml"), "r")
+            call(os.path.join(under_test.TEST_SUITE_DIR, "suite_name.yml"), "r")
             for _ in range(len(suites))
         ]
         m.assert_has_calls(calls, any_order=True)
-        filename = os.path.join(grt.CONFIG_DIR, "suite_name_{{:0{}}}.yml".format(
+        filename = os.path.join(under_test.CONFIG_DIR, "suite_name_{{:0{}}}.yml".format(
             int(math.ceil(math.log10(size)))))
         calls = [call(filename.format(i), "w") for i in range(size)]
         m.assert_has_calls(calls, any_order=True)
@@ -271,7 +352,7 @@ class RenderMiscSuites(unittest.TestCase):
         test_list = ["test{}".format(i) for i in range(10)]
         m = mock_open(read_data=yaml.dump({"selector": {"roots": []}}))
         with patch(ns("open"), m, create=True):
-            render_misc_suite(test_list, "suite_name")
+            under_test.render_misc_suite(test_list, "suite_name")
         handle = m()
 
         # The other writes are for the headers.
@@ -290,9 +371,9 @@ class RenderMiscSuites(unittest.TestCase):
   - test9
   roots: []
 """)
-        calls = [call(os.path.join(grt.TEST_SUITE_DIR, "suite_name.yml"), "r")]
+        calls = [call(os.path.join(under_test.TEST_SUITE_DIR, "suite_name.yml"), "r")]
         m.assert_has_calls(calls, any_order=True)
-        filename = os.path.join(grt.CONFIG_DIR, "suite_name_misc.yml")
+        filename = os.path.join(under_test.CONFIG_DIR, "suite_name_misc.yml")
         calls = [call(filename, "w")]
         m.assert_has_calls(calls, any_order=True)
 
@@ -301,25 +382,25 @@ class PrepareDirectoryForSuite(unittest.TestCase):
     def test_no_directory(self):
         with patch(ns("os")) as mock_os:
             mock_os.path.exists.return_value = False
-            prepare_directory_for_suite("tmp")
+            under_test.prepare_directory_for_suite("tmp")
 
         mock_os.makedirs.assert_called_once_with("tmp")
 
 
 class CalculateTimeoutTest(unittest.TestCase):
     def test_min_timeout(self):
-        self.assertEqual(grt.MIN_TIMEOUT_SECONDS, grt.calculate_timeout(15, 1))
+        self.assertEqual(under_test.MIN_TIMEOUT_SECONDS, under_test.calculate_timeout(15, 1))
 
     def test_over_timeout_by_one_minute(self):
-        self.assertEqual(360, grt.calculate_timeout(301, 1))
+        self.assertEqual(360, under_test.calculate_timeout(301, 1))
 
     def test_float_runtimes(self):
-        self.assertEqual(360, grt.calculate_timeout(300.14, 1))
+        self.assertEqual(360, under_test.calculate_timeout(300.14, 1))
 
     def test_scaling_factor(self):
         scaling_factor = 10
-        self.assertEqual(grt.MIN_TIMEOUT_SECONDS * scaling_factor,
-                         grt.calculate_timeout(30, scaling_factor))
+        self.assertEqual(under_test.MIN_TIMEOUT_SECONDS * scaling_factor,
+                         under_test.calculate_timeout(30, scaling_factor))
 
 
 class EvergreenConfigGeneratorTest(unittest.TestCase):
@@ -327,7 +408,7 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
     def generate_mock_suites(count):
         suites = []
         for idx in range(count):
-            suite = Mock()
+            suite = MagicMock()
             suite.name = "suite {0}".format(idx)
             suite.max_runtime = 5.28
             suite.get_runtime = lambda: 100.874
@@ -337,7 +418,7 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
 
     @staticmethod
     def generate_mock_options():
-        options = Mock()
+        options = MagicMock()
         options.resmoke_args = "resmoke_args"
         options.run_multiple_jobs = "true"
         options.variant = "buildvariant"
@@ -351,11 +432,18 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
 
         return options
 
+    @staticmethod
+    def mock_task(display_name):
+        mock = MagicMock()
+        mock.display_name = display_name
+        return mock
+
     def test_evg_config_is_created(self):
         options = self.generate_mock_options()
         suites = self.generate_mock_suites(3)
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         self.assertEqual(len(config["tasks"]), len(suites) + 1)
         command1 = config["tasks"][0]["commands"][2]
@@ -368,7 +456,8 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         options.task = "task"
         suites = self.generate_mock_suites(3)
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         self.assertEqual(len(config["tasks"]), len(suites) + 1)
         display_task = config["buildvariants"][0]["display_tasks"][0]
@@ -389,14 +478,15 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
 
         suites = self.generate_mock_suites(3)
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         self.assertEqual(len(config["tasks"]), len(suites) + 1)
         self.assertEqual(options.large_distro_name,
                          config["buildvariants"][0]["tasks"][0]["distros"][0])
 
     def test_selecting_tasks(self):
-        is_task_dependency = grt.EvergreenConfigGenerator._is_task_dependency
+        is_task_dependency = under_test.EvergreenConfigGenerator._is_task_dependency
         self.assertFalse(is_task_dependency("sharding", "sharding"))
         self.assertFalse(is_task_dependency("sharding", "other_task"))
         self.assertFalse(is_task_dependency("sharding", "sharding_gen"))
@@ -409,15 +499,15 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         options = self.generate_mock_options()
         suites = self.generate_mock_suites(3)
 
-        cfg_generator = grt.EvergreenConfigGenerator(suites, options, Mock())
+        cfg_generator = under_test.EvergreenConfigGenerator(suites, options, MagicMock())
         cfg_generator.build_tasks = [
-            {"display_name": "sharding_gen"},
-            {"display_name": "sharding_0"},
-            {"display_name": "other_task"},
-            {"display_name": "other_task_2"},
-            {"display_name": "sharding_1"},
-            {"display_name": "compile"},
-            {"display_name": "sharding_misc"},
+            self.mock_task("sharding_gen"),
+            self.mock_task("sharding_0"),
+            self.mock_task("other_task"),
+            self.mock_task("other_task_2"),
+            self.mock_task("sharding_1"),
+            self.mock_task("compile"),
+            self.mock_task("sharding_misc"),
         ]
 
         dependent_tasks = cfg_generator._get_tasks_for_depends_on("sharding")
@@ -432,18 +522,18 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         options.is_patch = False
         suites = self.generate_mock_suites(3)
 
-        cfg_generator = grt.EvergreenConfigGenerator(suites, options, Mock())
+        cfg_generator = under_test.EvergreenConfigGenerator(suites, options, MagicMock())
         cfg_generator.build_tasks = [
-            {"display_name": "sharding_gen"},
-            {"display_name": "sharding_0"},
-            {"display_name": "other_task"},
-            {"display_name": "other_task_2"},
-            {"display_name": "sharding_1"},
-            {"display_name": "compile"},
-            {"display_name": "sharding_misc"},
+            self.mock_task("sharding_gen"),
+            self.mock_task("sharding_0"),
+            self.mock_task("other_task"),
+            self.mock_task("other_task_2"),
+            self.mock_task("sharding_1"),
+            self.mock_task("compile"),
+            self.mock_task("sharding_misc"),
         ]
 
-        cfg_mock = Mock()
+        cfg_mock = MagicMock()
         cfg_generator._add_dependencies(cfg_mock)
         self.assertEqual(4, cfg_mock.dependency.call_count)
 
@@ -452,7 +542,8 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         options.repeat_suites = 5
         suites = self.generate_mock_suites(3)
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         self.assertEqual(len(config["tasks"]), len(suites) + 1)
         command1 = config["tasks"][0]["commands"][2]
@@ -460,18 +551,41 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         self.assertIn(options.resmoke_args, command1["vars"]["resmoke_args"])
         timeout_cmd = config["tasks"][0]["commands"][0]
         self.assertEqual("timeout.update", timeout_cmd["command"])
-        expected_timeout = grt.calculate_timeout(suites[0].max_runtime, 3) * 5
+        expected_timeout = under_test.calculate_timeout(suites[0].max_runtime, 3) * 5
         self.assertEqual(expected_timeout, timeout_cmd["params"]["timeout_secs"])
-        expected_exec_timeout = grt.calculate_timeout(suites[0].get_runtime(), 3) * 5
+        expected_exec_timeout = under_test.calculate_timeout(suites[0].get_runtime(), 3) * 5
         self.assertEqual(expected_exec_timeout, timeout_cmd["params"]["exec_timeout_secs"])
 
+    def test_evg_config_does_not_overwrite_repeatSuites_resmoke_arg_with_repeatSuites_default(self):
+        options = self.generate_mock_options()
+        options.resmoke_args = "resmoke_args --repeatSuites=5"
+        suites = self.generate_mock_suites(1)
+
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
+        command1 = config["tasks"][0]["commands"][2]
+        self.assertIn("--repeatSuites=5", command1["vars"]["resmoke_args"])
+        self.assertNotIn("--repeatSuites=1", command1["vars"]["resmoke_args"])
+
+    def test_evg_config_does_not_overwrite_repeat_resmoke_arg_with_repeatSuites_default(self):
+        options = self.generate_mock_options()
+        options.resmoke_args = "resmoke_args --repeat=5"
+        suites = self.generate_mock_suites(1)
+
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
+        command1 = config["tasks"][0]["commands"][2]
+        self.assertIn("--repeat=5", command1["vars"]["resmoke_args"])
+        self.assertNotIn("--repeatSuites=1", command1["vars"]["resmoke_args"])
+
     def test_suites_without_enough_info_should_not_include_timeouts(self):
         suite_without_timing_info = 1
         options = self.generate_mock_options()
         suites = self.generate_mock_suites(3)
         suites[suite_without_timing_info].should_overwrite_timeout.return_value = False
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         timeout_cmd = config["tasks"][suite_without_timing_info]["commands"][0]
         self.assertNotIn("command", timeout_cmd)
@@ -483,7 +597,8 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
         suites = self.generate_mock_suites(3)
         options.use_default_timeouts = True
 
-        config = grt.EvergreenConfigGenerator(suites, options, Mock()).generate_config().to_map()
+        config = under_test.EvergreenConfigGenerator(suites, options,
+                                                     MagicMock()).generate_config().to_map()
 
         timeout_cmd = config["tasks"][suite_without_timing_info]["commands"][0]
         self.assertNotIn("command", timeout_cmd)
@@ -492,35 +607,46 @@ class EvergreenConfigGeneratorTest(unittest.TestCase):
 
 class NormalizeTestNameTest(unittest.TestCase):
     def test_unix_names(self):
-        self.assertEqual("/home/user/test.js", grt.normalize_test_name("/home/user/test.js"))
+        self.assertEqual("/home/user/test.js", under_test.normalize_test_name("/home/user/test.js"))
 
     def test_windows_names(self):
-        self.assertEqual("/home/user/test.js", grt.normalize_test_name("\\home\\user\\test.js"))
+        self.assertEqual("/home/user/test.js",
+                         under_test.normalize_test_name("\\home\\user\\test.js"))
 
 
-class MainTest(unittest.TestCase):
+class GenerateSubSuitesTest(unittest.TestCase):
     @staticmethod
     def get_mock_options():
-        options = Mock()
+        options = MagicMock()
         options.target_resmoke_time = 10
         options.fallback_num_sub_suites = 2
         return options
 
+    @staticmethod
+    def get_test_list(n_tests):
+        return ["test{i}.js".format(i=i) for i in range(n_tests)]
+
+    @staticmethod
+    def history_mock(i):
+        mock = MagicMock()
+        mock.test_file = "test{}.js".format(i)
+        mock.avg_duration_pass = 60
+        mock.num_pass = 1
+        return mock
+
     def test_calculate_suites(self):
-        evg = Mock()
-        evg.test_stats.return_value = [{
-            "test_file": "test{}.js".format(i), "avg_duration_pass": 60, "num_pass": 1
-        } for i in range(100)]
+        evg = MagicMock()
+        evg.test_stats_by_project.return_value = [self.history_mock(i) for i in range(100)]
+        config_options = self.get_mock_options()
+        config_options.max_sub_suites = 1000
 
-        main = grt.Main(evg)
-        main.options = Mock()
-        main.config_options = self.get_mock_options()
+        gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
 
         with patch("os.path.exists") as exists_mock, patch(ns("suitesconfig")) as suitesconfig_mock:
             exists_mock.return_value = True
             suitesconfig_mock.get_suite.return_value.tests = \
-                [stat["test_file"] for stat in evg.test_stats.return_value]
-            suites = main.calculate_suites(_DATE, _DATE)
+                [stat.test_file for stat in evg.test_stats_by_project.return_value]
+            suites = gen_sub_suites.calculate_suites(_DATE, _DATE)
 
             # There are 100 tests taking 1 minute, with a target of 10 min we expect 10 suites.
             self.assertEqual(10, len(suites))
@@ -529,77 +655,72 @@ class MainTest(unittest.TestCase):
 
     def test_calculate_suites_fallback(self):
         n_tests = 100
-        response = Mock()
+        response = MagicMock()
         response.status_code = requests.codes.SERVICE_UNAVAILABLE
-        evg = Mock()
-        evg.test_stats.side_effect = requests.HTTPError(response=response)
+        evg = MagicMock()
+        evg.test_stats_by_project.side_effect = requests.HTTPError(response=response)
+        config_options = self.get_mock_options()
 
-        main = grt.Main(evg)
-        main.options = Mock()
-        main.options.execution_time_minutes = 10
-        main.config_options = self.get_mock_options()
-        main.list_tests = Mock(return_value=["test{}.js".format(i) for i in range(n_tests)])
+        gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+        gen_sub_suites.list_tests = MagicMock(return_value=self.get_test_list(n_tests))
 
-        suites = main.calculate_suites(_DATE, _DATE)
+        suites = gen_sub_suites.calculate_suites(_DATE, _DATE)
 
-        self.assertEqual(main.config_options.fallback_num_sub_suites, len(suites))
+        self.assertEqual(gen_sub_suites.config_options.fallback_num_sub_suites, len(suites))
         for suite in suites:
             self.assertEqual(50, len(suite.tests))
 
-        self.assertEqual(n_tests, len(main.test_list))
+        self.assertEqual(n_tests, len(gen_sub_suites.test_list))
 
     def test_calculate_suites_uses_fallback_for_no_results(self):
         n_tests = 100
-        evg = Mock()
-        evg.test_stats.return_value = []
+        evg = MagicMock()
+        evg.test_stats_by_project.return_value = []
+        config_options = self.get_mock_options()
 
-        main = grt.Main(evg)
-        main.options = Mock()
-        main.config_options = self.get_mock_options()
-        main.list_tests = Mock(return_value=["test{}.js".format(i) for i in range(n_tests)])
-        suites = main.calculate_suites(_DATE, _DATE)
+        gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+        gen_sub_suites.list_tests = MagicMock(return_value=self.get_test_list(n_tests))
+        suites = gen_sub_suites.calculate_suites(_DATE, _DATE)
 
-        self.assertEqual(main.config_options.fallback_num_sub_suites, len(suites))
+        self.assertEqual(gen_sub_suites.config_options.fallback_num_sub_suites, len(suites))
         for suite in suites:
             self.assertEqual(50, len(suite.tests))
 
-        self.assertEqual(n_tests, len(main.test_list))
+        self.assertEqual(n_tests, len(gen_sub_suites.test_list))
 
     def test_calculate_suites_uses_fallback_if_only_results_are_filtered(self):
         n_tests = 100
-        evg = Mock()
-        evg.test_stats.return_value = [{
-            "test_file": "test{}.js".format(i), "avg_duration_pass": 60, "num_pass": 1
-        } for i in range(100)]
-
-        main = grt.Main(evg)
-        main.options = Mock()
-        main.config_options = self.get_mock_options()
-        main.list_tests = Mock(return_value=["test{}.js".format(i) for i in range(n_tests)])
+        evg = MagicMock()
+        evg.test_stats_by_project.return_value = [
+            MagicMock(test_file="test{}.js".format(i), avg_duration_pass=60, num_pass=1)
+            for i in range(100)
+        ]
+        config_options = self.get_mock_options()
+
+        gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+        gen_sub_suites.list_tests = MagicMock(return_value=self.get_test_list(n_tests))
         with patch("os.path.exists") as exists_mock:
             exists_mock.return_value = False
-            suites = main.calculate_suites(_DATE, _DATE)
+            suites = gen_sub_suites.calculate_suites(_DATE, _DATE)
 
-            self.assertEqual(main.config_options.fallback_num_sub_suites, len(suites))
+            self.assertEqual(gen_sub_suites.config_options.fallback_num_sub_suites, len(suites))
             for suite in suites:
                 self.assertEqual(50, len(suite.tests))
 
-            self.assertEqual(n_tests, len(main.test_list))
+            self.assertEqual(n_tests, len(gen_sub_suites.test_list))
 
     def test_calculate_suites_error(self):
-        response = Mock()
+        response = MagicMock()
         response.status_code = requests.codes.INTERNAL_SERVER_ERROR
-        evg = Mock()
-        evg.test_stats.side_effect = requests.HTTPError(response=response)
+        evg = MagicMock()
+        evg.test_stats_by_project.side_effect = requests.HTTPError(response=response)
+        config_options = self.get_mock_options()
 
-        main = grt.Main(evg)
-        main.options = Mock()
-        main.options.execution_time_minutes = 10
-        main.config_options = self.get_mock_options()
-        main.list_tests = Mock(return_value=["test{}.js".format(i) for i in range(100)])
+        gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+        gen_sub_suites.list_tests = MagicMock(return_value=self.get_test_list(100))
 
         with self.assertRaises(requests.HTTPError):
-            main.calculate_suites(_DATE, _DATE)
+            gen_sub_suites.calculate_suites(_DATE, _DATE)
 
     def test_filter_missing_files(self):
         tests_runtimes = [
@@ -610,13 +731,13 @@ class MainTest(unittest.TestCase):
 
         with patch("os.path.exists") as exists_mock, patch(ns("suitesconfig")) as suitesconfig_mock:
             exists_mock.side_effect = [False, True, True]
-            evg = Mock()
+            evg = MagicMock()
             suitesconfig_mock.get_suite.return_value.tests = \
                 [runtime[0] for runtime in tests_runtimes]
-            main = grt.Main(evg)
-            main.config_options = Mock()
-            main.config_options.suite = "suite"
-            filtered_list = main.filter_existing_tests(tests_runtimes)
+            config_options = MagicMock(suite="suite")
+
+            gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+            filtered_list = gen_sub_suites.filter_existing_tests(tests_runtimes)
 
             self.assertEqual(2, len(filtered_list))
             self.assertNotIn(tests_runtimes[0], filtered_list)
@@ -634,13 +755,13 @@ class MainTest(unittest.TestCase):
 
         with patch("os.path.exists") as exists_mock, patch(ns("suitesconfig")) as suitesconfig_mock:
             exists_mock.return_value = True
-            evg = Mock()
+            evg = MagicMock()
             suitesconfig_mock.get_suite.return_value.tests = \
                 [runtime[0] for runtime in tests_runtimes if runtime[0] != blacklisted_test]
-            main = grt.Main(evg)
-            main.config_options = Mock()
-            main.config_options.suite = "suite"
-            filtered_list = main.filter_existing_tests(tests_runtimes)
+            config_options = MagicMock(suite="suite")
+
+            gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+            filtered_list = gen_sub_suites.filter_existing_tests(tests_runtimes)
 
             self.assertEqual(2, len(filtered_list))
             self.assertNotIn(blacklisted_test, filtered_list)
@@ -658,15 +779,15 @@ class MainTest(unittest.TestCase):
 
         with patch("os.path.exists") as exists_mock, patch(ns("suitesconfig")) as suitesconfig_mock:
             exists_mock.return_value = True
-            evg = Mock()
+            evg = MagicMock()
             suitesconfig_mock.get_suite.return_value.tests = [
                 runtime[0].replace("/", "\\") for runtime in tests_runtimes
                 if runtime[0] != blacklisted_test
             ]
-            main = grt.Main(evg)
-            main.config_options = Mock()
-            main.config_options.suite = "suite"
-            filtered_list = main.filter_existing_tests(tests_runtimes)
+            config_options = MagicMock(suite="suite")
+
+            gen_sub_suites = under_test.GenerateSubSuites(evg, config_options)
+            filtered_list = gen_sub_suites.filter_existing_tests(tests_runtimes)
 
             self.assertNotIn(blacklisted_test, filtered_list)
             self.assertIn(tests_runtimes[2], filtered_list)