SERVER-60458 Add a resmoke run flag to limit number of tests being run

7 files changed, 127 insertions, 76 deletions

diff --git a/buildscripts/resmokelib/config.py b/buildscripts/resmokelib/config.py
index 21b37a91f35..7141e206602 100644
--- a/buildscripts/resmokelib/config.py
+++ b/buildscripts/resmokelib/config.py
@@ -155,6 +155,9 @@ DEFAULTS = {
 
     # Generate multiversion exclude tags options
     "exclude_tags_file_path": "generated_resmoke_config/multiversion_exclude_tags.yml",
+
+    # Limit the number of tests to execute
+    "max_test_queue_size": None,
 }
 
 _SuiteOptions = collections.namedtuple("_SuiteOptions", [
@@ -540,6 +543,9 @@ UNDO_RECORDER_PATH = None
 # # Generate multiversion exclude tags options
 EXCLUDE_TAGS_FILE_PATH = None
 
+# Limit the number of tests to execute
+MAX_TEST_QUEUE_SIZE = None
+
 ##
 # Internally used configuration options that aren't exposed to the user
 ##
diff --git a/buildscripts/resmokelib/configure_resmoke.py b/buildscripts/resmokelib/configure_resmoke.py
index 50da8c1b0f0..3b624fdcbc9 100644
--- a/buildscripts/resmokelib/configure_resmoke.py
+++ b/buildscripts/resmokelib/configure_resmoke.py
@@ -368,6 +368,8 @@ def _update_config_vars(values):  # pylint: disable=too-many-statements,too-many
 
     _config.EXCLUDE_TAGS_FILE_PATH = config.pop("exclude_tags_file_path")
 
+    _config.MAX_TEST_QUEUE_SIZE = config.pop("max_test_queue_size")
+
     def configure_tests(test_files, replay_file):
         # `_validate_options` has asserted that at most one of `test_files` and `replay_file` contains input.
 
diff --git a/buildscripts/resmokelib/run/__init__.py b/buildscripts/resmokelib/run/__init__.py
index 303b4a267d6..279185791e6 100644
--- a/buildscripts/resmokelib/run/__init__.py
+++ b/buildscripts/resmokelib/run/__init__.py
@@ -843,6 +843,9 @@ class RunPlugin(PluginInterface):
                             action="append", metavar="featureFlag1, featureFlag2, ...",
                             help="Additional feature flags")
 
+        parser.add_argument("--maxTestQueueSize", type=int, dest="max_test_queue_size",
+                            help=argparse.SUPPRESS)
+
         mongodb_server_options = parser.add_argument_group(
             title=_MONGODB_SERVER_OPTIONS_TITLE,
             description=("Options related to starting a MongoDB cluster that are forwarded from"
diff --git a/buildscripts/resmokelib/testing/executor.py b/buildscripts/resmokelib/testing/executor.py
index 65f04dd18a2..504e3d460c8 100644
--- a/buildscripts/resmokelib/testing/executor.py
+++ b/buildscripts/resmokelib/testing/executor.py
@@ -2,6 +2,7 @@
 
 import threading
 import time
+from typing import List
 
 from buildscripts.resmokelib import config as _config
 from buildscripts.resmokelib import errors
@@ -14,8 +15,8 @@ from buildscripts.resmokelib.testing import hooks as _hooks
 from buildscripts.resmokelib.testing import job as _job
 from buildscripts.resmokelib.testing import report as _report
 from buildscripts.resmokelib.testing import testcases
-from buildscripts.resmokelib.testing.queue_element import queue_elem_factory
-from buildscripts.resmokelib.utils.queue import Queue
+from buildscripts.resmokelib.testing.queue_element import queue_elem_factory, QueueElem
+from buildscripts.resmokelib.utils import queue as _queue
 
 
 class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
@@ -52,11 +53,8 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
                                                       archive)
 
         self._suite = suite
-        self.num_tests = len(suite.tests) * suite.options.num_repeat_tests
         self.test_queue_logger = logging.loggers.new_testqueue_logger(suite.test_kind)
-
-        # Must be done after getting buildlogger configuration.
-        self._jobs = self._create_jobs(self.num_tests)
+        self._jobs = []
 
     def _num_jobs_to_start(self, suite, num_tests):
         """
@@ -109,6 +107,7 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
             num_repeat_suites = self._suite.options.num_repeat_suites
             while num_repeat_suites > 0:
                 test_queue = self._make_test_queue()
+                self._jobs = self._create_jobs(test_queue.num_tests)
 
                 partial_reports = [job.report for job in self._jobs]
                 self._suite.record_test_start(partial_reports)
@@ -157,10 +156,10 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
                 test_report = report.as_dict()
                 test_results_num = len(test_report["results"])
                 # There should be at least as many tests results as expected number of tests.
-                if test_results_num < self.num_tests:
+                if test_results_num < test_queue.num_tests:
                     raise errors.ResmokeError(
                         "{} reported tests is less than {} expected tests".format(
-                            test_results_num, self.num_tests))
+                            test_results_num, test_queue.num_tests))
 
                 # Clear the report so it can be reused for the next execution.
                 for job in self._jobs:
@@ -230,7 +229,7 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
         # We cannot return 'interrupt_flag.is_set()' because the interrupt flag can be set by a Job
         # instance if a test fails and it decides to drain the queue. We only want to raise a
         # StopExecution exception in TestSuiteExecutor.run() if the user triggered the interrupt.
-        return (combined_report, user_interrupted)
+        return combined_report, user_interrupted
 
     def _teardown_fixtures(self):
         """Tear down all of the fixtures.
@@ -291,16 +290,6 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
         return _job.Job(job_num, job_logger, fixture, hooks, report, self.archival,
                         self._suite.options, self.test_queue_logger)
 
-    def _num_times_to_repeat_tests(self):
-        """
-        Determine the number of times to repeat the tests.
-
-        :return: Number of times to repeat the tests.
-        """
-        if self._suite.options.num_repeat_tests:
-            return self._suite.options.num_repeat_tests
-        return 1
-
     def _create_queue_elem_for_test_name(self, test_name):
         """
         Create the appropriate queue_elem to run the given test_name.
@@ -321,24 +310,47 @@ class TestSuiteExecutor(object):  # pylint: disable=too-many-instance-attributes
         we will add 2 queue_elements of each test to the queue). If we are repeating execution for
         a specified time period, we will add each test to the queue, but as a QueueElemRepeatTime
         object, which will requeue itself if it has not run for the expected duration.
-
         Use a multi-consumer queue instead of a unittest.TestSuite so that the test cases can
         be dispatched to multiple threads.
-
         :return: Queue of testcases to run.
         """
-        queue = Queue()
+        test_queue = TestQueue()
 
-        # Put all the test cases in a queue.
-        for _ in range(self._num_times_to_repeat_tests()):
+        # Make test cases to put in test queue
+        test_cases = []
+        for _ in range(self._suite.get_num_times_to_repeat_tests()):
             for test_name in self._suite.tests:
                 queue_elem = self._create_queue_elem_for_test_name(test_name)
-                queue.put(queue_elem)
+                test_cases.append(queue_elem)
+        test_queue.add_test_cases(test_cases)
 
-        return queue
+        return test_queue
 
     def _log_timeout_warning(self, seconds):
         """Log a message if any thread fails to terminate after `seconds`."""
         self.logger.warning(
             '*** Still waiting for processes to terminate after %s seconds. Try using ctrl-\\ '
             'to send a SIGQUIT on Linux or ctrl-c again on Windows ***', seconds)
+
+
+class TestQueue(_queue.Queue):
+    """A queue of test cases to run.
+
+    Use a multi-consumer queue instead of a unittest.TestSuite so that the test cases can
+    be dispatched to multiple threads.
+    """
+
+    def __init__(self):
+        """Initialize test queue."""
+        self.num_tests = 0
+        self.max_test_queue_size = utils.default_if_none(_config.MAX_TEST_QUEUE_SIZE, -1)
+        super().__init__()
+
+    def add_test_cases(self, test_cases: List[QueueElem]) -> None:
+        """Add test cases to the queue."""
+        for test_case in test_cases:
+            if self.max_test_queue_size < 0 or self.num_tests < self.max_test_queue_size:
+                self.put(test_case)
+                self.num_tests += 1
+            else:
+                break
diff --git a/buildscripts/resmokelib/testing/suite.py b/buildscripts/resmokelib/testing/suite.py
index 2910406eaf1..5de8bdd5855 100644
--- a/buildscripts/resmokelib/testing/suite.py
+++ b/buildscripts/resmokelib/testing/suite.py
@@ -166,6 +166,12 @@ class Suite(object):  # pylint: disable=too-many-instance-attributes
         """Return the "test_kind" section of the YAML configuration."""
         return self._suite_config["test_kind"]
 
+    def get_num_times_to_repeat_tests(self) -> int:
+        """Return the number of times to repeat tests."""
+        if self.options.num_repeat_tests:
+            return self.options.num_repeat_tests
+        return 1
+
     @property
     def options(self):
         """Get the options."""
diff --git a/buildscripts/tests/resmokelib/testing/test_executor.py b/buildscripts/tests/resmokelib/testing/test_executor.py
index e4a95fcdbac..275503f11c4 100644
--- a/buildscripts/tests/resmokelib/testing/test_executor.py
+++ b/buildscripts/tests/resmokelib/testing/test_executor.py
@@ -5,7 +5,6 @@ import unittest
 import mock
 
 from buildscripts.resmokelib.testing import executor
-from buildscripts.resmokelib.testing import queue_element
 
 # pylint: disable=missing-docstring,protected-access
 
@@ -21,36 +20,16 @@ def mock_suite(n_tests):
     suite = mock.MagicMock()
     suite.test_kind = "js_test"
     suite.tests = ["jstests/core/and{}.js".format(i) for i in range(n_tests)]
-    suite.options.num_repeat_tests = None
+    suite.get_num_times_to_repeat_tests.return_value = 1
     return suite
 
 
-class TestTestSuiteExecutor(unittest.TestCase):
-    def test__make_test_queue_time_repeat(self):
-        suite = mock_suite(2)
-        suite.options.time_repeat_tests_secs = 30
-        executor_object = UnitTestExecutor(suite, {})
-        test_queue = executor_object._make_test_queue()
-        self.assertFalse(test_queue.empty())
-        self.assertEqual(test_queue.qsize(), len(suite.tests))
-        for suite_test in suite.tests:
-            test_element = test_queue.get_nowait()
-            self.assertIsInstance(test_element, queue_element.QueueElemRepeatTime)
-            self.assertEqual(test_element.testcase.test_name, suite_test)
-        self.assertTrue(test_queue.empty())
-
-    def test__make_test_queue_num_repeat(self):
-        suite = mock_suite(2)
-        suite.options.time_repeat_tests_secs = None
-        executor_object = UnitTestExecutor(suite, {})
-        test_queue = executor_object._make_test_queue()
-        self.assertFalse(test_queue.empty())
-        self.assertEqual(test_queue.qsize(), len(suite.tests))
-        for suite_test in suite.tests:
-            test_element = test_queue.get_nowait()
-            self.assertIsInstance(test_element, queue_element.QueueElem)
-            self.assertEqual(test_element.testcase.test_name, suite_test)
-        self.assertTrue(test_queue.empty())
+class UnitTestExecutor(executor.TestSuiteExecutor):
+    def __init__(self, suite, config):  # pylint: disable=super-init-not-called
+        self._suite = suite
+        self.test_queue_logger = logging.getLogger("executor_unittest")
+        self.test_config = config
+        self.logger = mock.MagicMock()
 
 
 class TestNumJobsToStart(unittest.TestCase):
@@ -91,21 +70,6 @@ class TestCreateJobs(unittest.TestCase):
         self.assertEqual(num_jobs, self.ut_executor._make_job.call_count)
 
 
-class TestNumTimesToRepeatTests(unittest.TestCase):
-    def test_default(self):
-        num_tests = 1
-        suite = mock_suite(num_tests)
-        ut_executor = UnitTestExecutor(suite, None)
-        self.assertEqual(1, ut_executor._num_times_to_repeat_tests())
-
-    def test_with_num_repeat_tests(self):
-        num_tests = 1
-        suite = mock_suite(num_tests)
-        suite.options.num_repeat_tests = 5
-        ut_executor = UnitTestExecutor(suite, None)
-        self.assertEqual(suite.options.num_repeat_tests, ut_executor._num_times_to_repeat_tests())
-
-
 class TestCreateQueueElemForTestName(unittest.TestCase):
     @mock.patch(ns("testcases.make_test_case"))
     @mock.patch(ns("queue_elem_factory"))
@@ -137,7 +101,7 @@ class TestMakeTestQueue(unittest.TestCase):
 
     def test_repeat_three_times(self):
         num_repeats = 3
-        self.suite.options.num_repeat_tests = num_repeats
+        self.suite.get_num_times_to_repeat_tests.return_value = num_repeats
         test_queue = self.ut_executor._make_test_queue()
         self.assertEqual(num_repeats * len(self.suite.tests), test_queue.qsize())
         while not test_queue.empty():
@@ -145,9 +109,41 @@ class TestMakeTestQueue(unittest.TestCase):
             self.assertIn(element, self.suite.tests)
 
 
-class UnitTestExecutor(executor.TestSuiteExecutor):
-    def __init__(self, suite, config):  # pylint: disable=super-init-not-called
-        self._suite = suite
-        self.test_queue_logger = logging.getLogger("executor_unittest")
-        self.test_config = config
-        self.logger = mock.MagicMock()
+class TestTestQueueAddTestCases(unittest.TestCase):
+    def setUp(self):
+        self.default_max_test_queue_size = executor._config.MAX_TEST_QUEUE_SIZE
+        self.num_test_cases = 3
+        self.test_cases = [mock.MagicMock() for _ in range(self.num_test_cases)]
+
+    def tearDown(self):
+        executor._config.MAX_TEST_QUEUE_SIZE = self.default_max_test_queue_size
+
+    def test_do_not_set_max_test_queue_size(self):
+        test_queue = executor.TestQueue()
+        test_queue.add_test_cases(self.test_cases)
+        self.assertEqual(test_queue.num_tests, self.num_test_cases)
+        while not test_queue.empty():
+            element = test_queue.get()
+            self.assertIn(element, self.test_cases)
+
+    def test_max_test_queue_size_not_reached(self):
+        max_test_queue_size = 10
+        self.assertTrue(max_test_queue_size > self.num_test_cases)
+        executor._config.MAX_TEST_QUEUE_SIZE = max_test_queue_size
+        test_queue = executor.TestQueue()
+        test_queue.add_test_cases(self.test_cases)
+        self.assertEqual(test_queue.num_tests, self.num_test_cases)
+        while not test_queue.empty():
+            element = test_queue.get()
+            self.assertIn(element, self.test_cases)
+
+    def test_max_test_queue_size_exceeded(self):
+        max_test_queue_size = 2
+        self.assertTrue(max_test_queue_size < self.num_test_cases)
+        executor._config.MAX_TEST_QUEUE_SIZE = max_test_queue_size
+        test_queue = executor.TestQueue()
+        test_queue.add_test_cases(self.test_cases)
+        self.assertEqual(test_queue.num_tests, max_test_queue_size)
+        while not test_queue.empty():
+            element = test_queue.get()
+            self.assertIn(element, self.test_cases)
diff --git a/buildscripts/tests/resmokelib/testing/test_suite.py b/buildscripts/tests/resmokelib/testing/test_suite.py
new file mode 100644
index 00000000000..5ed95757818
--- /dev/null
+++ b/buildscripts/tests/resmokelib/testing/test_suite.py
@@ -0,0 +1,26 @@
+"""Unit tests for the resmokelib.testing.suite module."""
+import unittest
+
+from buildscripts.resmokelib.testing import suite as under_test
+
+# pylint: disable=missing-docstring,protected-access
+
+
+class TestNumTimesToRepeatTests(unittest.TestCase):
+    def setUp(self):
+        self.default_repeat_tests = under_test._config.REPEAT_TESTS
+        self.suite = under_test.Suite("suite_name", {"test_kind": "js_test"})
+
+    def tearDown(self):
+        under_test._config.REPEAT_TESTS = self.default_repeat_tests
+
+    def test_without_num_repeat_tests(self):
+        expected_num_repeat_tests = 1
+        num_repeat_tests = self.suite.get_num_times_to_repeat_tests()
+        self.assertEqual(num_repeat_tests, expected_num_repeat_tests)
+
+    def test_with_num_repeat_tests(self):
+        expected_num_repeat_tests = 5
+        under_test._config.REPEAT_TESTS = expected_num_repeat_tests
+        num_repeat_tests = self.suite.get_num_times_to_repeat_tests()
+        self.assertEqual(num_repeat_tests, expected_num_repeat_tests)