SERVER-38111: Add max_sub_suites in generate_resmoke_suite

(cherry picked from commit b12875f9e31321c247cc70515d669dd71ffbf700)

3 files changed, 81 insertions, 5 deletions

diff --git a/buildscripts/generate_resmoke_suites.py b/buildscripts/generate_resmoke_suites.py
index 6564af7cdd8..6268de9f258 100755
--- a/buildscripts/generate_resmoke_suites.py
+++ b/buildscripts/generate_resmoke_suites.py
@@ -201,7 +201,19 @@ def sort_list_of_test_by_max_runtime(tests):
     return sorted(tests.keys(), key=lambda test: tests[test][MAX_RUNTIME_KEY], reverse=True)
 
 
-def divide_tests_into_suites_by_maxtime(tests, sorted_tests, max_time_seconds):
+def divide_remaining_tests_among_suites(remaining_tests, tests, suites):
+    """Divide the list of tests given among the suites given."""
+    suite_idx = 0
+    for test_name in remaining_tests:
+        test = tests[test_name]
+        current_suite = suites[suite_idx]
+        current_suite.add_test(test_name, test)
+        suite_idx += 1
+        if suite_idx >= len(suites):
+            suite_idx = 0
+
+
+def divide_tests_into_suites_by_maxtime(tests, sorted_tests, max_time_seconds, max_suites=None):
     """
     Divide the given tests into suites.
 
@@ -209,8 +221,9 @@ def divide_tests_into_suites_by_maxtime(tests, sorted_tests, max_time_seconds):
     """
     suites = []
     current_suite = Suite()
+    last_test_processed = len(sorted_tests)
     LOGGER.debug("Determines suites for runtime: %ds", max_time_seconds)
-    for test_name in sorted_tests:
+    for idx, test_name in enumerate(sorted_tests):
         test = tests[test_name]
         if current_suite.get_runtime() + test[MAX_RUNTIME_KEY] > max_time_seconds:
             LOGGER.debug("Runtime(%d) + new test(%d) > max(%d)", current_suite.get_runtime(),
@@ -218,12 +231,19 @@ def divide_tests_into_suites_by_maxtime(tests, sorted_tests, max_time_seconds):
             if current_suite.get_test_count() > 0:
                 suites.append(current_suite)
                 current_suite = Suite()
+                if max_suites and len(suites) >= max_suites:
+                    last_test_processed = idx
+                    break
 
         current_suite.add_test(test_name, test)
 
     if current_suite.get_test_count() > 0:
         suites.append(current_suite)
 
+    if max_suites and last_test_processed < len(sorted_tests):
+        # We must have hit the max suite limit, just randomly add the remaining tests to suites.
+        divide_remaining_tests_among_suites(sorted_tests[last_test_processed:], tests, suites)
+
     return suites
 
 
@@ -328,6 +348,8 @@ class Main(object):
         parser.add_argument("--variants", dest="variants", metavar="<variant1,variant2,...>",
                             default=None,
                             help="Comma-separated list of Evergreeen build variants to analyze.")
+        parser.add_argument("--max-sub-suites", dest="max_sub_suites", type=int,
+                            help="Max number of suites to divide into.")
         parser.add_argument("--verbose", dest="verbose", action="store_true", default=False,
                             help="Enable verbose logging.")
         parser.add_argument("task", nargs=1, help="task to analyze.")
@@ -351,7 +373,8 @@ class Main(object):
         """Divide test into suites that can be run in less than the specified execution time."""
         tests = organize_executions_by_test(data)
         self.test_list = sort_list_of_test_by_max_runtime(tests)
-        return divide_tests_into_suites_by_maxtime(tests, self.test_list, execution_time_secs)
+        return divide_tests_into_suites_by_maxtime(tests, self.test_list, execution_time_secs,
+                                                   self.options.max_sub_suites)
 
     def render_suites(self, suites, task):
         """Render the given suites into yml files that can be used by resmoke.py."""
diff --git a/buildscripts/resmokeconfig/suites/buildscripts_test.yml b/buildscripts/resmokeconfig/suites/buildscripts_test.yml
index 8c6093b7853..4886acb7b43 100644
--- a/buildscripts/resmokeconfig/suites/buildscripts_test.yml
+++ b/buildscripts/resmokeconfig/suites/buildscripts_test.yml
@@ -7,9 +7,7 @@ selector:
   exclude_files:
   - buildscripts/tests/resmokelib/test_archival.py # Requires boto3.
   - buildscripts/tests/resmokelib/test_selector.py # Test assumes POSIX path.
-  - buildscripts/tests/resmokelib/testing/hooks/test_combine_benchmark_results.py # Requires mock.
   - buildscripts/tests/test_aws_ec2.py # Requires boto3.
-  - buildscripts/tests/test_generate_resmoke_suites.py # Requires mock.
   - buildscripts/tests/test_evergreen_gen_fuzzer_tests.py # Requires mock.
   - buildscripts/tests/test_remote_operations.py # Requires ssh to be enabled locally.
   - buildscripts/tests/test_update_test_lifecycle.py # Test assumes POSIX path.
diff --git a/buildscripts/tests/test_generate_resmoke_suites.py b/buildscripts/tests/test_generate_resmoke_suites.py
index 024c3a582f5..efb69f061ed 100644
--- a/buildscripts/tests/test_generate_resmoke_suites.py
+++ b/buildscripts/tests/test_generate_resmoke_suites.py
@@ -398,6 +398,41 @@ class OrganizeExecutionsByTestTest(unittest.TestCase):
         self.assertEquals(len(tests), 0)
 
 
+class DivideRemainingTestsAmongSuitesTest(unittest.TestCase):
+    @staticmethod
+    def generate_tests(n_tests):
+        tests = {}
+        test_names = []
+        for idx in range(n_tests):
+            name = "test_{0}".format(idx)
+            test_names.append(name)
+            tests[name] = {"max_runtime": 2 * idx}
+
+        return test_names, tests
+
+    def test_each_suite_gets_one_test(self):
+        suites = [grs.Suite(), grs.Suite(), grs.Suite()]
+        test_names, tests = self.generate_tests(3)
+
+        grs.divide_remaining_tests_among_suites(test_names, tests, suites)
+
+        for suite in suites:
+            self.assertEqual(suite.get_test_count(), 1)
+
+    def test_each_suite_gets_at_least_one_test(self):
+        suites = [grs.Suite(), grs.Suite(), grs.Suite()]
+        test_names, tests = self.generate_tests(5)
+
+        grs.divide_remaining_tests_among_suites(test_names, tests, suites)
+
+        total_tests = 0
+        for suite in suites:
+            total_tests += suite.get_test_count()
+            self.assertGreaterEqual(suite.get_test_count(), 1)
+
+        self.assertEqual(total_tests, len(tests))
+
+
 class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
     def test_if_less_total_than_max_only_one_suite_created(self):
         max_time = 20
@@ -439,6 +474,26 @@ class DivideTestsIntoSuitesByMaxtimeTest(unittest.TestCase):
         self.assertEqual(suites[0].get_test_count(), 1)
         self.assertEqual(suites[0].get_runtime(), 15)
 
+    def test_max_sub_suites_options(self):
+        max_time = 5
+        max_suites = 2
+        test_names = ["test1", "test2", "test3", "test4", "test5"]
+        tests = {
+            test_names[0]: {"max_runtime": 5},
+            test_names[1]: {"max_runtime": 4},
+            test_names[2]: {"max_runtime": 3},
+            test_names[3]: {"max_runtime": 4},
+            test_names[4]: {"max_runtime": 3},
+        }
+
+        suites = grs.divide_tests_into_suites_by_maxtime(tests, test_names, max_time,
+                                                         max_suites=max_suites)
+        self.assertEqual(len(suites), max_suites)
+        total_tests = 0
+        for suite in suites:
+            total_tests += suite.get_test_count()
+        self.assertEqual(total_tests, len(test_names))
+
 
 class SuiteTest(unittest.TestCase):
     def test_adding_tests_increases_count_and_runtime(self):