SERVER-46236 Selected_tests_gen task should run tasks on all required variants, not just enterprise-rhel-62-64-bit

2 files changed, 162 insertions, 117 deletions

diff --git a/buildscripts/selected_tests.py b/buildscripts/selected_tests.py
index 3a37fb98ea8..f9464e0e30f 100644
--- a/buildscripts/selected_tests.py
+++ b/buildscripts/selected_tests.py
@@ -89,23 +89,22 @@ class SelectedTestsConfigOptions(ConfigOptions):
 
     @classmethod
     # pylint: disable=too-many-arguments,W0221
-    def from_file(cls, origin_variant_expansions: Dict[str, str], filepath: str,
-                  overwrites: Dict[str, Any], required_keys: Set[str], defaults: Dict[str, Any],
-                  formats: Dict[str, type]):
+    def from_file(cls, origin_variant_expansions: Dict[str, str],
+                  selected_tests_variant_expansions: Dict[str, str], overwrites: Dict[str, Any],
+                  required_keys: Set[str], defaults: Dict[str, Any], formats: Dict[str, type]):
         """
         Create an instance of SelectedTestsConfigOptions based on the given config file.
 
         :param origin_variant_expansions: Expansions of the origin build variant.
-        :param filepath: Path to file containing configuration.
+        :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
         :param overwrites: Dict of configuration values to overwrite those listed in filepath.
         :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 SelectedTestsConfigOptions.
         """
-        config_from_file = read_config.read_config_file(filepath)
-        return cls({**origin_variant_expansions, **config_from_file, **overwrites}, required_keys,
-                   defaults, formats)
+        return cls({**origin_variant_expansions, **selected_tests_variant_expansions, **overwrites},
+                   required_keys, defaults, formats)
 
     @property
     def run_tests_task(self):
@@ -214,30 +213,30 @@ def _find_task(build_variant_config: Variant, task_name: str) -> Task:
     return task
 
 
-def _get_selected_tests_task_config(expansion_file: str) -> Dict[str, str]:
+def _get_selected_tests_task_config(
+        selected_tests_variant_expansions: Dict[str, str]) -> Dict[str, str]:
     """
     Look up task config of the selected tests task.
 
-    :param expansion_file: Configuration file.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
     :return: Task configuration values.
     """
-    expansions = read_config.read_config_file(expansion_file)
     return {
-        "name_of_generating_task": expansions["task_name"],
-        "name_of_generating_build_variant": expansions["build_variant"],
-        "name_of_generating_build_id": expansions["build_id"]
+        "name_of_generating_task": selected_tests_variant_expansions["task_name"],
+        "name_of_generating_build_variant": selected_tests_variant_expansions["build_variant"],
+        "name_of_generating_build_id": selected_tests_variant_expansions["build_id"]
     }
 
 
 def _get_evg_task_config(
-        expansion_file: str,
+        selected_tests_variant_expansions: Dict[str, str],
         task_name: str,
         build_variant_config: Variant,
 ) -> Dict[str, Any]:
     """
     Look up task config of the task to be generated.
 
-    :param expansion_file: Configuration file.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
     :param task_name: Task to get info for.
     :param build_variant_config: Config of build variant to collect task info from.
     :return: Task configuration values.
@@ -258,7 +257,7 @@ def _get_evg_task_config(
     resmoke_args_without_suites = ResmokeArgs.remove_arg(task_vars["resmoke_args"], "suites")
     task_vars["resmoke_args"] = resmoke_args_without_suites
 
-    selected_tests_task_config = _get_selected_tests_task_config(expansion_file)
+    selected_tests_task_config = _get_selected_tests_task_config(selected_tests_variant_expansions)
 
     return {
         "task_name": task.name, "build_variant": build_variant_config.name, **task_vars,
@@ -286,19 +285,21 @@ def _update_config_with_task(evg_api: EvergreenApi, shrub_config: Configuration,
     task_generator.generate_task_config(shrub_config, suites)
 
 
-def _get_task_configs_for_test_mappings(expansion_file: str, tests_by_task: Dict[str, Any],
+def _get_task_configs_for_test_mappings(selected_tests_variant_expansions: Dict[str, str],
+                                        tests_by_task: Dict[str, Any],
                                         build_variant_config: Variant) -> Dict[str, dict]:
     """
     For test mappings, generate a dict containing task names and their config settings.
 
-    :param expansion_file: Configuration file.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
     :param tests_by_task: Dictionary of tests and tasks to run.
     :param build_variant_config: Config of build variant to collect task info from.
     :return: Dict of task names and their config settings.
     """
     evg_task_configs = {}
     for task_name, test_list_info in tests_by_task.items():
-        evg_task_config = _get_evg_task_config(expansion_file, task_name, build_variant_config)
+        evg_task_config = _get_evg_task_config(selected_tests_variant_expansions, task_name,
+                                               build_variant_config)
         evg_task_config.update({"selected_tests_to_run": set(test_list_info["tests"])})
         LOGGER.debug("Calculated evg_task_config values", evg_task_config=evg_task_config)
         evg_task_configs[task_name] = evg_task_config
@@ -306,54 +307,52 @@ def _get_task_configs_for_test_mappings(expansion_file: str, tests_by_task: Dict
     return evg_task_configs
 
 
-def _get_task_configs_for_task_mappings(expansion_file: str, related_tasks: List[str],
+def _get_task_configs_for_task_mappings(selected_tests_variant_expansions: Dict[str, str],
+                                        related_tasks: List[str],
                                         build_variant_config: Variant) -> Dict[str, dict]:
     """
     For task mappings, generate a dict containing task names and their config settings.
 
-    :param expansion_file: Configuration file.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
     :param related_tasks: List of tasks to run.
     :param build_variant_config: Config of build variant to collect task info from.
     :return: Dict of task names and their config settings.
     """
     evg_task_configs = {}
     for task_name in related_tasks:
-        evg_task_config = _get_evg_task_config(expansion_file, task_name, build_variant_config)
+        evg_task_config = _get_evg_task_config(selected_tests_variant_expansions, task_name,
+                                               build_variant_config)
         LOGGER.debug("Calculated evg_task_config values", evg_task_config=evg_task_config)
         evg_task_configs[task_name] = evg_task_config
 
     return evg_task_configs
 
 
-# pylint: disable=too-many-arguments, too-many-locals
-def run(evg_api: EvergreenApi, evg_conf: EvergreenProjectConfig, expansion_file: str,
-        selected_tests_service: SelectedTestsService, changed_files: Set[str],
-        build_variant: str) -> Dict[str, dict]:
+def _get_task_configs(evg_conf: EvergreenProjectConfig,
+                      selected_tests_service: SelectedTestsService,
+                      selected_tests_variant_expansions: Dict[str, str],
+                      build_variant_config: Variant, changed_files: Set[str]) -> Dict[str, Dict]:
     """
-    Run code to select tasks to run based on test mappings and task mappings.
+    Get task configurations for the tasks to be generated.
 
-    :param evg_api: Evergreen API object.
     :param evg_conf: Evergreen configuration.
-    :param expansion_file: Configuration file.
-    :param selected_tests_config: Location of config file to connect to elected-tests service.
+    :param selected_tests_service: Selected-tests service.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
+    :param build_variant_config: Config of build variant to collect task info from.
     :param changed_files: Set of changed_files.
-    :param build_variant: Build variant to collect task info from.
-    :return: Dict of files and file contents for generated tasks.
+    :return: Task configurations.
     """
-    shrub_config = Configuration()
-    config_dict_of_suites_and_tasks = {}
-
     task_configs = {}
-    build_variant_config = evg_conf.get_variant(build_variant)
 
     related_test_files = _find_selected_test_files(selected_tests_service, changed_files)
     LOGGER.debug("related test files found", related_test_files=related_test_files)
     if related_test_files:
-        tests_by_task = create_task_list_for_tests(related_test_files, build_variant, evg_conf)
+        tests_by_task = create_task_list_for_tests(related_test_files, build_variant_config.name,
+                                                   evg_conf)
         LOGGER.debug("tests and tasks found", tests_by_task=tests_by_task)
 
         test_mapping_task_configs = _get_task_configs_for_test_mappings(
-            expansion_file, tests_by_task, build_variant_config)
+            selected_tests_variant_expansions, tests_by_task, build_variant_config)
         task_configs.update(test_mapping_task_configs)
 
     related_tasks = _find_selected_tasks(selected_tests_service, changed_files,
@@ -361,25 +360,53 @@ def run(evg_api: EvergreenApi, evg_conf: EvergreenProjectConfig, expansion_file:
     LOGGER.debug("related tasks found", related_tasks=related_tasks)
     if related_tasks:
         task_mapping_task_configs = _get_task_configs_for_task_mappings(
-            expansion_file, related_tasks, build_variant_config)
+            selected_tests_variant_expansions, related_tasks, build_variant_config)
         # task_mapping_task_configs will overwrite test_mapping_task_configs
         # because task_mapping_task_configs will run all tests rather than a subset of tests and we
         # should err on the side of running all tests
         task_configs.update(task_mapping_task_configs)
 
-    origin_variant_expansions = build_variant_config.expansions
+    return task_configs
 
-    for task_config in task_configs.values():
-        config_options = SelectedTestsConfigOptions.from_file(
-            origin_variant_expansions,
-            expansion_file,
-            task_config,
-            REQUIRED_CONFIG_KEYS,
-            DEFAULT_CONFIG_VALUES,
-            CONFIG_FORMAT_FN,
-        )
-        _update_config_with_task(evg_api, shrub_config, config_options,
-                                 config_dict_of_suites_and_tasks)
+
+# pylint: disable=too-many-arguments
+def run(evg_api: EvergreenApi, evg_conf: EvergreenProjectConfig,
+        selected_tests_service: SelectedTestsService,
+        selected_tests_variant_expansions: Dict[str, str], changed_files: Set[str],
+        origin_build_variants: List[str]) -> Dict[str, dict]:
+    """
+    Run code to select tasks to run based on test and task mappings for each of the build variants.
+
+    :param evg_api: Evergreen API object.
+    :param evg_conf: Evergreen configuration.
+    :param selected_tests_service: Selected-tests service.
+    :param selected_tests_variant_expansions: Expansions of the selected-tests variant.
+    :param changed_files: Set of changed_files.
+    :param origin_build_variants: Build variants to collect task info from.
+    :return: Dict of files and file contents for generated tasks.
+    """
+    shrub_config = Configuration()
+    config_dict_of_suites_and_tasks = {}
+
+    for build_variant in origin_build_variants:
+        build_variant_config = evg_conf.get_variant(build_variant)
+        origin_variant_expansions = build_variant_config.expansions
+
+        task_configs = _get_task_configs(evg_conf, selected_tests_service,
+                                         selected_tests_variant_expansions, build_variant_config,
+                                         changed_files)
+
+        for task_config in task_configs.values():
+            config_options = SelectedTestsConfigOptions.from_file(
+                origin_variant_expansions,
+                selected_tests_variant_expansions,
+                task_config,
+                REQUIRED_CONFIG_KEYS,
+                DEFAULT_CONFIG_VALUES,
+                CONFIG_FORMAT_FN,
+            )
+            _update_config_with_task(evg_api, shrub_config, config_options,
+                                     config_dict_of_suites_and_tasks)
 
     config_dict_of_suites_and_tasks["selected_tests_config.json"] = shrub_config.to_json()
     return config_dict_of_suites_and_tasks
@@ -402,12 +429,6 @@ def run(evg_api: EvergreenApi, evg_conf: EvergreenProjectConfig, expansion_file:
     help="Configuration file with connection info for Evergreen API.",
 )
 @click.option(
-    "--build-variant",
-    "build_variant",
-    required=True,
-    help="Tasks to run will be selected from this build variant.",
-)
-@click.option(
     "--selected-tests-config",
     "selected_tests_config",
     required=True,
@@ -418,7 +439,6 @@ def main(
         verbose: bool,
         expansion_file: str,
         evg_api_config: str,
-        build_variant: str,
         selected_tests_config: str,
 ):
     """
@@ -427,7 +447,6 @@ def main(
     :param verbose: Log extra debug information.
     :param expansion_file: Configuration file.
     :param evg_api_config: Location of configuration file to connect to evergreen.
-    :param build_variant: Build variant to query tasks from.
     :param selected_tests_config: Location of config file to connect to elected-tests service.
     """
     _configure_logging(verbose)
@@ -441,8 +460,13 @@ def main(
     buildscripts.resmokelib.parser.set_options()
     LOGGER.debug("Found changed files", files=changed_files)
 
-    config_dict_of_suites_and_tasks = run(evg_api, evg_conf, expansion_file, selected_tests_service,
-                                          changed_files, build_variant)
+    selected_tests_variant_expansions = read_config.read_config_file(expansion_file)
+    origin_build_variants = selected_tests_variant_expansions["selected_tests_buildvariants"].split(
+        " ")
+
+    config_dict_of_suites_and_tasks = run(evg_api, evg_conf, selected_tests_service,
+                                          selected_tests_variant_expansions, changed_files,
+                                          origin_build_variants)
     write_file_dict(SELECTED_TESTS_CONFIG_DIR, config_dict_of_suites_and_tasks)
 
 
diff --git a/buildscripts/tests/test_selected_tests.py b/buildscripts/tests/test_selected_tests.py
index 6d88e1e8f49..d8a4e0bc349 100644
--- a/buildscripts/tests/test_selected_tests.py
+++ b/buildscripts/tests/test_selected_tests.py
@@ -38,51 +38,47 @@ def tests_by_task_stub():
 
 
 class TestSelectedTestsConfigOptions(unittest.TestCase):
-    @patch(ns("read_config"))
-    def test_overwrites_overwrite_filepath_config(self, read_config_mock):
-        filepath = MagicMock()
+    def test_overwrites_overwrite_filepath_config(self):
         origin_variant_expansions = {"key1": 0}
-        read_config_mock.read_config_file.return_value = {"key1": 1}
+        selected_tests_variant_expansions = {"key1": 1}
         overwrites = {"key1": 2}
         required_keys = {"key1"}
         defaults = {}
         formats = {"key1": int}
 
         config_options = under_test.SelectedTestsConfigOptions.from_file(
-            origin_variant_expansions, filepath, overwrites, required_keys, defaults, formats)
+            origin_variant_expansions, selected_tests_variant_expansions, overwrites, required_keys,
+            defaults, formats)
 
         self.assertEqual(overwrites["key1"], config_options.key1)
 
-    @patch(ns("read_config"))
-    def test_overwrites_overwrite_defaults(self, read_config_mock):
-        filepath = MagicMock()
+    def test_overwrites_overwrite_defaults(self):
         origin_variant_expansions = {}
-        read_config_mock.read_config_file.return_value = {"key1": 1}
+        selected_tests_variant_expansions = {"key1": 1}
         overwrites = {"key1": 2}
         required_keys = {"key1"}
         defaults = {"key1": 3}
         formats = {"key1": int}
 
         config_options = under_test.SelectedTestsConfigOptions.from_file(
-            origin_variant_expansions, filepath, overwrites, required_keys, defaults, formats)
+            origin_variant_expansions, selected_tests_variant_expansions, overwrites, required_keys,
+            defaults, formats)
 
         self.assertEqual(overwrites["key1"], config_options.key1)
 
-    @patch(ns("read_config"))
-    def test_filepath_config_overrides_origin_expansions(self, read_config_mock):
-        filepath = MagicMock()
+    def test_selected_tests_config_overrides_origin_expansions(self):
         origin_variant_expansions = {"key1": 0}
-        filepath_config = {"key1": 1}
-        read_config_mock.read_config_file.return_value = filepath_config
+        selected_tests_variant_expansions = {"key1": 1}
         overwrites = {}
         required_keys = {"key1"}
         defaults = {}
         formats = {"key1": int}
 
         config_options = under_test.SelectedTestsConfigOptions.from_file(
-            origin_variant_expansions, filepath, overwrites, required_keys, defaults, formats)
+            origin_variant_expansions, selected_tests_variant_expansions, overwrites, required_keys,
+            defaults, formats)
 
-        self.assertEqual(filepath_config["key1"], config_options.key1)
+        self.assertEqual(selected_tests_variant_expansions["key1"], config_options.key1)
 
     def test_run_tests_task(self):
         config_options = under_test.SelectedTestsConfigOptions(
@@ -259,13 +255,13 @@ class TestFindSelectedTasks(unittest.TestCase):
 
 
 class TestGetSelectedTestsTaskConfiguration(unittest.TestCase):
-    @patch(ns("read_config"))
-    def test_gets_values(self, read_config_mock):
-        read_config_mock.read_config_file.return_value = {
+    def test_gets_values(self):
+        selected_tests_variant_expansions = {
             "task_name": "my_task", "build_variant": "my-build-variant", "build_id": "my_build_id"
         }
 
-        selected_tests_task_config = under_test._get_selected_tests_task_config(MagicMock())
+        selected_tests_task_config = under_test._get_selected_tests_task_config(
+            selected_tests_variant_expansions)
 
         self.assertEqual(selected_tests_task_config["name_of_generating_task"], "my_task")
         self.assertEqual(selected_tests_task_config["name_of_generating_build_variant"],
@@ -292,8 +288,7 @@ class TestGetEvgTaskConfig(unittest.TestCase):
             }],
         })
 
-        evg_task_config = under_test._get_evg_task_config(MagicMock(), task_name,
-                                                          build_variant_conf)
+        evg_task_config = under_test._get_evg_task_config({}, task_name, build_variant_conf)
 
         self.assertEqual(evg_task_config["task_name"], task_name)
         self.assertEqual(evg_task_config["build_variant"], "variant")
@@ -319,8 +314,7 @@ class TestGetEvgTaskConfig(unittest.TestCase):
             }],
         })
 
-        evg_task_config = under_test._get_evg_task_config(MagicMock(), task_name,
-                                                          build_variant_conf)
+        evg_task_config = under_test._get_evg_task_config({}, task_name, build_variant_conf)
 
         self.assertEqual(evg_task_config["task_name"], task_name)
         self.assertEqual(evg_task_config["build_variant"], "variant")
@@ -382,9 +376,8 @@ class TestGetTaskConfigsForTestMappings(unittest.TestCase):
         get_evg_task_config_mock.side_effect = [{"task_config_key": "task_config_value_1"},
                                                 {"task_config_key": "task_config_value_2"}]
 
-        task_configs = under_test._get_task_configs_for_test_mappings(
-            expansion_file=MagicMock(), tests_by_task=tests_by_task,
-            build_variant_config=MagicMock())
+        task_configs = under_test._get_task_configs_for_test_mappings({}, tests_by_task,
+                                                                      MagicMock())
 
         self.assertEqual(task_configs["jsCore_auth"]["task_config_key"], "task_config_value_1")
         self.assertEqual(
@@ -401,67 +394,95 @@ class TestGetTaskConfigsForTaskMappings(unittest.TestCase):
         tasks = ["task_1", "task_2"]
         get_evg_task_config_mock.side_effect = [{"task_config_key": "task_config_value_1"},
                                                 {"task_config_key": "task_config_value_2"}]
-        task_configs = under_test._get_task_configs_for_task_mappings(
-            expansion_file=MagicMock(), related_tasks=tasks, build_variant_config=MagicMock())
+        task_configs = under_test._get_task_configs_for_task_mappings({}, tasks, MagicMock())
 
         self.assertEqual(task_configs["task_1"]["task_config_key"], "task_config_value_1")
         self.assertEqual(task_configs["task_2"]["task_config_key"], "task_config_value_2")
 
 
-class TestRun(unittest.TestCase):
-    @patch(ns("SelectedTestsConfigOptions"))
+class TestGetTaskConfigs(unittest.TestCase):
     @patch(ns("_find_selected_test_files"))
     @patch(ns("create_task_list_for_tests"))
     @patch(ns("_get_task_configs_for_test_mappings"))
     @patch(ns("_find_selected_tasks"))
-    @patch(ns("_update_config_with_task"))
-    # pylint: disable=too-many-arguments
-    def test_run_with_related_tests_but_no_related_tasks(
-            self, update_config_with_task_mock, find_selected_tasks_mock,
-            get_task_configs_for_test_mappings_mock, create_task_list_for_tests_mock,
-            find_selected_test_files_mock, selected_tests_config_options):
+    def test_with_related_tests_but_no_related_tasks(
+            self, find_selected_tasks_mock, get_task_configs_for_test_mappings_mock,
+            create_task_list_for_tests_mock, find_selected_test_files_mock):
         find_selected_test_files_mock.return_value = {"jstests/file-1.js", "jstests/file-3.js"}
         get_task_configs_for_test_mappings_mock.return_value = {
             "task_config_key": "task_config_value_1"
         }
         find_selected_tasks_mock.return_value = set()
+        changed_files = {"src/file1.cpp", "src/file2.js"}
 
-        def update_config_with_task(evg_api, shrub_config, config_options,
-                                    config_dict_of_suites_and_tasks):
-            config_dict_of_suites_and_tasks["new_config_key"] = "new_config_values"
+        task_configs = under_test._get_task_configs(MagicMock(), MagicMock(), {}, MagicMock(),
+                                                    changed_files)
 
-        update_config_with_task_mock.side_effect = update_config_with_task
+        self.assertEqual(task_configs["task_config_key"], "task_config_value_1")
 
+    @patch(ns("_find_selected_test_files"))
+    @patch(ns("create_task_list_for_tests"))
+    @patch(ns("_get_task_configs_for_task_mappings"))
+    @patch(ns("_find_selected_tasks"))
+    def test_with_no_related_tests_but_related_tasks(
+            self, find_selected_tasks_mock, get_task_configs_for_task_mappings_mock,
+            create_task_list_for_tests_mock, find_selected_test_files_mock):
+        find_selected_test_files_mock.return_value = {}
+        find_selected_tasks_mock.return_value = {"jsCore_auth", "auth_gen"}
+        get_task_configs_for_task_mappings_mock.return_value = {
+            "task_config_key": "task_config_value_2"
+        }
         changed_files = {"src/file1.cpp", "src/file2.js"}
-        config_dict_of_suites_and_tasks = under_test.run(MagicMock(), MagicMock(), MagicMock(),
-                                                         MagicMock(), changed_files, "variant")
 
-        self.assertEqual(config_dict_of_suites_and_tasks["new_config_key"], "new_config_values")
+        task_configs = under_test._get_task_configs(MagicMock(), MagicMock(), {}, MagicMock(),
+                                                    changed_files)
+
+        self.assertEqual(task_configs["task_config_key"], "task_config_value_2")
 
-    @patch(ns("SelectedTestsConfigOptions"))
     @patch(ns("_find_selected_test_files"))
+    @patch(ns("create_task_list_for_tests"))
+    @patch(ns("_get_task_configs_for_test_mappings"))
     @patch(ns("_get_task_configs_for_task_mappings"))
     @patch(ns("_find_selected_tasks"))
-    @patch(ns("_update_config_with_task"))
     # pylint: disable=too-many-arguments
-    def test_run_with_related_tasks_but_no_related_tests(
-            self, update_config_with_task_mock, find_selected_tasks_mock,
-            get_task_configs_for_task_mappings_mock, find_selected_test_files_mock,
-            selected_tests_config_options):
-        find_selected_test_files_mock.return_value = {}
-        get_task_configs_for_task_mappings_mock.return_value = {
+    def test_task_mapping_configs_will_overwrite_test_mapping_configs(
+            self, find_selected_tasks_mock, get_task_configs_for_task_mappings_mock,
+            get_task_configs_for_test_mappings_mock, create_task_list_for_tests_mock,
+            find_selected_test_files_mock):
+        find_selected_test_files_mock.return_value = {"jstests/file-1.js", "jstests/file-3.js"}
+        get_task_configs_for_test_mappings_mock.return_value = {
             "task_config_key": "task_config_value_1"
         }
         find_selected_tasks_mock.return_value = {"jsCore_auth", "auth_gen"}
+        get_task_configs_for_task_mappings_mock.return_value = {
+            "task_config_key": "task_config_value_2"
+        }
+        changed_files = {"src/file1.cpp", "src/file2.js"}
+
+        task_configs = under_test._get_task_configs(MagicMock(), MagicMock(), {}, MagicMock(),
+                                                    changed_files)
+
+        self.assertEqual(task_configs["task_config_key"], "task_config_value_2")
+
+
+class TestRun(unittest.TestCase):
+    @patch(ns("SelectedTestsConfigOptions"))
+    @patch(ns("_get_task_configs"))
+    @patch(ns("_update_config_with_task"))
+    def test_run(self, update_config_with_task_mock, get_task_configs_mock,
+                 selected_tests_config_options):
+        get_task_configs_mock.return_value = {"task_config_key": "task_config_value_1"}
 
         def update_config_with_task(evg_api, shrub_config, config_options,
                                     config_dict_of_suites_and_tasks):
             config_dict_of_suites_and_tasks["new_config_key"] = "new_config_values"
+            shrub_config.task("my_fake_task")
 
         update_config_with_task_mock.side_effect = update_config_with_task
-
         changed_files = {"src/file1.cpp", "src/file2.js"}
-        config_dict_of_suites_and_tasks = under_test.run(MagicMock(), MagicMock(), MagicMock(),
-                                                         MagicMock(), changed_files, "variant")
+
+        config_dict_of_suites_and_tasks = under_test.run(MagicMock(), MagicMock(), MagicMock(), {},
+                                                         changed_files, ["variant_1", "variant_2"])
 
         self.assertEqual(config_dict_of_suites_and_tasks["new_config_key"], "new_config_values")
+        self.assertIn("my_fake_task", config_dict_of_suites_and_tasks["selected_tests_config.json"])