SERVER-42440: burn_in_tests should run tasks on the distro they normally run on

(cherry picked from commit 0ef81e97a8f60329c44640e6b76189a3b51bd42a)

2 files changed, 279 insertions, 114 deletions

diff --git a/buildscripts/burn_in_tests.py b/buildscripts/burn_in_tests.py
index 8e510c57ceb..347d4bebcd3 100644
--- a/buildscripts/burn_in_tests.py
+++ b/buildscripts/burn_in_tests.py
@@ -10,7 +10,7 @@ import os.path
 import subprocess
 import shlex
 import sys
-import urllib.parse
+from typing import Dict, List
 
 from git import Repo
 import structlog
@@ -31,6 +31,7 @@ if __name__ == "__main__" and __package__ is None:
 from buildscripts.patch_builds.change_data import find_changed_files
 from buildscripts import resmokelib
 from buildscripts.ciconfig import evergreen
+from buildscripts.ciconfig.evergreen import EvergreenProjectConfig, VariantTask
 # pylint: enable=wrong-import-position
 
 structlog.configure(logger_factory=LoggerFactory())
@@ -294,9 +295,57 @@ def _set_resmoke_args(task):
     return evergreen.ResmokeArgs.get_updated_arg(resmoke_args, "suites", suite_name)
 
 
-def create_task_list(  #pylint: disable=too-many-locals
-        evergreen_conf, buildvariant, suites, exclude_tasks):
-    """Find associated tasks for the specified buildvariant and suites.
+def _distro_to_run_task_on(task: VariantTask, evg_proj_config: EvergreenProjectConfig,
+                           build_variant: str) -> str:
+    """
+    Determine what distro an task should be run on.
+
+    For normal tasks, the distro will be the default for the build variant unless the task spec
+    specifies a particular distro to run on.
+
+    For generated tasks, the distro will be the default for the build variant unless (1) the
+    "use_large_distro" flag is set as a "var" in the "generate resmoke tasks" command of the
+    task definition and (2) the build variant defines the "large_distro_name" in its expansions.
+
+    :param task: Task being run.
+    :param evg_proj_config: Evergreen project configuration.
+    :param build_variant: Build Variant task is being run on.
+    :return: Distro task should be run on.
+    """
+    task_def = evg_proj_config.get_task(task.name)
+    if task_def.is_generate_resmoke_task:
+        resmoke_vars = task_def.generate_resmoke_tasks_command["vars"]
+        if "use_large_distro" in resmoke_vars:
+            bv = evg_proj_config.get_variant(build_variant)
+            if "large_distro_name" in bv.raw["expansions"]:
+                return bv.raw["expansions"]["large_distro_name"]
+
+    return task.run_on[0]
+
+
+def _gather_task_info(task: VariantTask, tests_by_suite: Dict,
+                      evg_proj_config: EvergreenProjectConfig, build_variant: str) -> Dict:
+    """
+    Gather the information needed to run the given task.
+
+    :param task: Task to be run.
+    :param tests_by_suite: Dict of suites.
+    :param evg_proj_config: Evergreen project configuration.
+    :param build_variant: Build variant task will be run on.
+    :return: Dictionary of information needed to run task.
+    """
+    return {
+        "resmoke_args": _set_resmoke_args(task),
+        "tests": tests_by_suite[task.resmoke_suite],
+        "use_multiversion": task.multiversion_path,
+        "distro": _distro_to_run_task_on(task, evg_proj_config, build_variant)
+    }  # yapf: disable
+
+
+def create_task_list(evergreen_conf: EvergreenProjectConfig, build_variant: str,
+                     tests_by_suite: Dict[str, List[str]], exclude_tasks: [str]):
+    """
+    Find associated tasks for the specified build_variant and suites.
 
     Returns a dict keyed by task_name, with executor, resmoke_args & tests, i.e.,
     {'jsCore_small_oplog':
@@ -304,28 +353,31 @@ def create_task_list(  #pylint: disable=too-many-locals
          'tests': ['jstests/core/all2.js', 'jstests/core/all3.js'],
          'use_multiversion': '/data/multiversion'}
     }
+
+    :param evergreen_conf: Evergreen configuration for project.
+    :param build_variant: Build variant to select tasks from.
+    :param tests_by_suite: Suites to be run.
+    :param exclude_tasks: Tasks to exclude.
+    :return: Dict of tasks to run with run configuration.
     """
 
-    evg_buildvariant = evergreen_conf.get_variant(buildvariant)
-    if not evg_buildvariant:
-        print("Buildvariant '{}' not found in {}".format(buildvariant, evergreen_conf.path))
-        sys.exit(1)
+    evg_build_variant = evergreen_conf.get_variant(build_variant)
+    if not evg_build_variant:
+        print("Buildvariant '{}' not found in {}".format(build_variant, evergreen_conf.path))
+        raise ValueError(f"Buildvariant ({build_variant} not found in evergreen configuration")
 
     # Find all the buildvariant tasks.
     exclude_tasks_set = set(exclude_tasks)
-    variant_task = {
+    all_variant_tasks = {
         _get_task_name(task): task
-        for task in evg_buildvariant.tasks
+        for task in evg_build_variant.tasks
         if task.name not in exclude_tasks_set and task.combined_resmoke_args
     }
 
     # Return the list of tasks to run for the specified suite.
     return {
-        task_name: {
-            "resmoke_args": _set_resmoke_args(task), "tests": suites[task.resmoke_suite],
-            "use_multiversion": task.multiversion_path
-        }
-        for task_name, task in variant_task.items() if task.resmoke_suite in suites
+        task_name: _gather_task_info(task, tests_by_suite, evergreen_conf, build_variant)
+        for task_name, task in all_variant_tasks.items() if task.resmoke_suite in tests_by_suite
     }
 
 
diff --git a/buildscripts/tests/test_burn_in_tests.py b/buildscripts/tests/test_burn_in_tests.py
index a9bbb03e73b..6f250b23e5b 100644
--- a/buildscripts/tests/test_burn_in_tests.py
+++ b/buildscripts/tests/test_burn_in_tests.py
@@ -765,114 +765,227 @@ class CreateExecutorList(unittest.TestCase):
         self.assertEqual(mock_suite_class.call_count, 0)
 
 
-class CreateTaskList(unittest.TestCase):
-    def test_create_task_list(self):
-        variant = "variantall"
-        suites = [SUITE1, SUITE2, SUITE3]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), len(VARIANTS["variantall"].tasks))
-        self.assertIn("task1", task_list)
-        self.assertEqual(task_list["task1"]["resmoke_args"], "--suites=suite1 var1arg1")
-        self.assertEqual(task_list["task1"]["tests"], SUITE1.tests)
-        self.assertIsNone(task_list["task1"]["use_multiversion"])
-        self.assertIn("task2", task_list)
-        self.assertEqual(task_list["task2"]["resmoke_args"], "--suites=suite1 var1arg2")
-        self.assertEqual(task_list["task2"]["tests"], SUITE1.tests)
-        self.assertIsNone(task_list["task2"]["use_multiversion"])
-        self.assertIn("task3", task_list)
-        self.assertEqual(task_list["task3"]["resmoke_args"], "--suites=suite1 var1arg3")
-        self.assertEqual(task_list["task3"]["tests"], SUITE1.tests)
-        self.assertIsNone(task_list["task3"]["use_multiversion"])
-
-    def test_create_task_list_multiversion(self):
-        variant = "variant_multiversion"
-        suites = [SUITE1, SUITE2, SUITE3]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), len(VARIANTS["variant_multiversion"].tasks))
-        self.assertEqual(task_list["multiversion_task"]["use_multiversion"], MULTIVERSION_PATH)
-
-    def test_create_task_list_gen_tasks(self):
-        variant = "variant_generate_tasks"
-        suites = [SUITE3]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), len(VARIANTS["variant_generate_tasks"].tasks))
-        self.assertIn(GENERATE_RESMOKE_TASKS_BASENAME, task_list)
-        self.assertEqual(task_list[GENERATE_RESMOKE_TASKS_BASENAME]["tests"], SUITE3.tests)
-        self.assertIsNone(task_list[GENERATE_RESMOKE_TASKS_BASENAME]["use_multiversion"])
-
-    def test_create_task_list_gen_tasks_multiversion(self):
-        variant = "variant_generate_tasks_multiversion"
-        suites = [SUITE3]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), len(VARIANTS["variant_generate_tasks_multiversion"].tasks))
-        self.assertEqual(task_list[GENERATE_RESMOKE_TASKS_BASENAME]["use_multiversion"],
-                         MULTIVERSION_PATH)
-
-    def test_create_task_list_gen_tasks_no_suite(self):
-        variant = "variant_generate_tasks_no_suite"
-        suites = [SUITE3]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), len(VARIANTS["variant_generate_tasks_no_suite"].tasks))
-        self.assertIn(GENERATE_RESMOKE_TASKS_BASENAME, task_list)
-        self.assertEqual(task_list[GENERATE_RESMOKE_TASKS_BASENAME]["tests"], SUITE3.tests)
-
+def create_variant_task_mock(task_name, suite_name, distro="distro"):
+    variant_task = MagicMock()
+    variant_task.name = task_name
+    variant_task.generated_task_name = task_name
+    variant_task.resmoke_suite = suite_name
+    variant_task.get_vars_suite_name.return_value = suite_name
+    variant_task.combined_resmoke_args = f"--suites={suite_name}"
+    variant_task.multiversion_path = None
+    variant_task.run_on = [distro]
+    return variant_task
+
+
+class TestGatherTaskInfo(unittest.TestCase):
+    def test_non_generated_task(self):
+        suite_name = "suite_1"
+        distro_name = "distro_1"
+        variant = "build_variant"
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_task.return_value.is_generate_resmoke_task = False
+        task_mock = create_variant_task_mock("task 1", suite_name, distro_name)
+        test_list = [f"test{i}.js" for i in range(3)]
+        tests_by_suite = {
+            suite_name: test_list,
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+
+        task_info = burn_in._gather_task_info(task_mock, tests_by_suite, evg_conf_mock, variant)
+
+        self.assertIn(suite_name, task_info["resmoke_args"])
+        for test in test_list:
+            self.assertIn(test, task_info["tests"])
+        self.assertIsNone(task_info["use_multiversion"])
+        self.assertEqual(distro_name, task_info["distro"])
+
+    def test_multiversion_task(self):
+        suite_name = "suite_1"
+        distro_name = "distro_1"
+        variant = "build_variant"
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_task.return_value.is_generate_resmoke_task = False
+        task_mock = create_variant_task_mock("task 1", suite_name, distro_name)
+        task_mock.multiversion_path = "/path/to/multiversion"
+        test_list = [f"test{i}.js" for i in range(3)]
+        tests_by_suite = {
+            suite_name: test_list,
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+
+        task_info = burn_in._gather_task_info(task_mock, tests_by_suite, evg_conf_mock, variant)
+
+        self.assertIn(suite_name, task_info["resmoke_args"])
+        for test in test_list:
+            self.assertIn(test, task_info["tests"])
+        self.assertEqual(task_mock.multiversion_path, task_info["use_multiversion"])
+        self.assertEqual(distro_name, task_info["distro"])
+
+    def test_generated_task_no_large_on_task(self):
+        suite_name = "suite_1"
+        distro_name = "distro_1"
+        variant = "build_variant"
+        evg_conf_mock = MagicMock()
+        task_def_mock = evg_conf_mock.get_task.return_value
+        task_def_mock.is_generate_resmoke_task = True
+        task_def_mock.generate_resmoke_tasks_command = {"vars": {}}
+        task_mock = create_variant_task_mock("task 1", suite_name, distro_name)
+        test_list = [f"test{i}.js" for i in range(3)]
+        tests_by_suite = {
+            suite_name: test_list,
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+
+        task_info = burn_in._gather_task_info(task_mock, tests_by_suite, evg_conf_mock, variant)
+
+        self.assertIn(suite_name, task_info["resmoke_args"])
+        for test in test_list:
+            self.assertIn(test, task_info["tests"])
+        self.assertIsNone(task_info["use_multiversion"])
+        self.assertEqual(distro_name, task_info["distro"])
+
+    def test_generated_task_no_large_on_build_variant(self):
+        suite_name = "suite_1"
+        distro_name = "distro_1"
+        variant = "build_variant"
+        evg_conf_mock = MagicMock()
+        task_def_mock = evg_conf_mock.get_task.return_value
+        task_def_mock.is_generate_resmoke_task = True
+        task_def_mock.generate_resmoke_tasks_command = {"vars": {"use_large_distro": True}}
+        task_mock = create_variant_task_mock("task 1", suite_name, distro_name)
+        test_list = [f"test{i}.js" for i in range(3)]
+        tests_by_suite = {
+            suite_name: test_list,
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+
+        task_info = burn_in._gather_task_info(task_mock, tests_by_suite, evg_conf_mock, variant)
+
+        self.assertIn(suite_name, task_info["resmoke_args"])
+        for test in test_list:
+            self.assertIn(test, task_info["tests"])
+        self.assertIsNone(task_info["use_multiversion"])
+        self.assertEqual(distro_name, task_info["distro"])
+
+    def test_generated_task_large_distro(self):
+        suite_name = "suite_1"
+        distro_name = "distro_1"
+        large_distro_name = "large_distro_1"
+        variant = "build_variant"
+        evg_conf_mock = MagicMock()
+        task_def_mock = evg_conf_mock.get_task.return_value
+        task_def_mock.is_generate_resmoke_task = True
+        task_def_mock.generate_resmoke_tasks_command = {"vars": {"use_large_distro": True}}
+        evg_conf_mock.get_variant.return_value.raw = {
+            "expansions": {
+                "large_distro_name": large_distro_name
+            }
+        }  # yapf: disable
+        task_mock = create_variant_task_mock("task 1", suite_name, distro_name)
+        test_list = [f"test{i}.js" for i in range(3)]
+        tests_by_suite = {
+            suite_name: test_list,
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+
+        task_info = burn_in._gather_task_info(task_mock, tests_by_suite, evg_conf_mock, variant)
+
+        self.assertIn(suite_name, task_info["resmoke_args"])
+        for test in test_list:
+            self.assertIn(test, task_info["tests"])
+        self.assertIsNone(task_info["use_multiversion"])
+        self.assertEqual(large_distro_name, task_info["distro"])
+
+
+class TestCreateTaskList(unittest.TestCase):
     def test_create_task_list_no_excludes(self):
-        variant = "variant1"
-        suites = [SUITE1, SUITE2]
-        exclude_suites = []
-        suite_list = _create_executor_list(suites, exclude_suites)
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), 1)
-        self.assertIn("task1", task_list)
-        self.assertEqual(task_list["task1"]["resmoke_args"], "--suites=suite1 var1arg1")
-        self.assertEqual(task_list["task1"]["tests"], SUITE1.tests)
-        self.assertNotIn("task2", task_list)
-        self.assertNotIn("task3", task_list)
+        variant = "variant name"
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_variant.return_value.tasks = [
+            create_variant_task_mock("task 1", "suite 1"),
+            create_variant_task_mock("task 2", "suite 2"),
+            create_variant_task_mock("task 3", "suite 3"),
+        ]
+        tests_by_suite = {
+            "suite 1": [f"test{i}.js" for i in range(3)],
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+        exclude_tasks = []
+
+        task_list = burn_in.create_task_list(evg_conf_mock, variant, tests_by_suite, exclude_tasks)
+
+        self.assertIn("task 1", task_list)
+        self.assertIn("task 2", task_list)
+        self.assertIn("task 3", task_list)
+        self.assertEqual(3, len(task_list))
+
+    def test_create_task_list_has_correct_task_info(self):
+        variant = "variant name"
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_variant.return_value.tasks = [
+            create_variant_task_mock("task 1", "suite_1", "distro 1"),
+        ]
+        evg_conf_mock.get_task.return_value.run_on = ["distro 1"]
+        tests_by_suite = {
+            "suite_1": [f"test{i}.js" for i in range(3)],
+        }
+        exclude_tasks = []
+
+        task_list = burn_in.create_task_list(evg_conf_mock, variant, tests_by_suite, exclude_tasks)
+
+        self.assertIn("task 1", task_list)
+        task_info = task_list["task 1"]
+        self.assertIn("suite_1", task_info["resmoke_args"])
+        for i in range(3):
+            self.assertIn(f"test{i}.js", task_info["tests"])
+        self.assertIsNone(task_info["use_multiversion"])
+        self.assertEqual("distro 1", task_info["distro"])
 
     def test_create_task_list_with_excludes(self):
-        variant = "variant2"
-        suites = [SUITE1, SUITE2, SUITE3]
-        suite_list = _create_executor_list(suites, [])
-        exclude_suites = ["suite2"]
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), 1)
-        self.assertIn("task3", task_list)
-        self.assertEqual(task_list["task3"]["resmoke_args"], "--suites=suite3 var2arg3")
-        self.assertEqual(task_list["task3"]["tests"], SUITE3.tests)
-        self.assertNotIn("task1", task_list)
-        self.assertNotIn("task2", task_list)
+        variant = "variant name"
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_variant.return_value.tasks = [
+            create_variant_task_mock("task 1", "suite 1"),
+            create_variant_task_mock("task 2", "suite 2"),
+            create_variant_task_mock("task 3", "suite 3"),
+        ]
+        tests_by_suite = {
+            "suite 1": [f"test{i}.js" for i in range(3)],
+            "suite 2": [f"test{i}.js" for i in range(1)],
+            "suite 3": [f"test{i}.js" for i in range(2)],
+        }
+        exclude_tasks = ["task 2"]
+
+        task_list = burn_in.create_task_list(evg_conf_mock, variant, tests_by_suite, exclude_tasks)
+
+        self.assertIn("task 1", task_list)
+        self.assertNotIn("task 2", task_list)
+        self.assertIn("task 3", task_list)
+        self.assertEqual(2, len(task_list))
 
     def test_create_task_list_no_suites(self):
-        variant = "variant2"
-        suite_list = {}
-        exclude_suites = ["suite2"]
-        task_list = burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, exclude_suites)
-        self.assertEqual(len(task_list), 0)
-        self.assertEqual(task_list, {})
+        variant = "variant name"
+        evg_conf_mock = MagicMock()
+        suite_dict = {}
 
-    def test_create_task_list_novariant(self):
-        class BadVariant(Exception):
-            pass
+        task_list = burn_in.create_task_list(evg_conf_mock, variant, suite_dict, [])
 
-        def _raise_bad_variant(code=0):
-            raise BadVariant("Bad variant {}".format(code))
+        self.assertEqual(task_list, {})
 
+    def test_build_variant_not_in_evg_project_config(self):
         variant = "novariant"
-        suites = [SUITE1, SUITE2, SUITE3]
-        suite_list = _create_executor_list(suites, [])
-        with patch("sys.exit", _raise_bad_variant):
-            with self.assertRaises(BadVariant):
-                burn_in.create_task_list(EVERGREEN_CONF, variant, suite_list, [])
+        evg_conf_mock = MagicMock()
+        evg_conf_mock.get_variant.return_value = None
+        suite_dict = {}
+
+        with self.assertRaises(ValueError):
+            burn_in.create_task_list(EVERGREEN_CONF, variant, suite_dict, [])
 
 
 class TestFindChangedTests(unittest.TestCase):