SERVER-53058: Refactor how historic test stats are collected

1 files changed, 12 insertions, 14 deletions

diff --git a/buildscripts/burn_in_tests.py b/buildscripts/burn_in_tests.py
index fe8780cb203..cb808a1f6a3 100644
--- a/buildscripts/burn_in_tests.py
+++ b/buildscripts/burn_in_tests.py
@@ -34,7 +34,7 @@ from buildscripts.resmokelib.utils import default_if_none, globstar
 from buildscripts.ciconfig.evergreen import parse_evergreen_file, ResmokeArgs, \
     EvergreenProjectConfig, VariantTask
 from buildscripts.util.fileops import write_file
-from buildscripts.util.teststats import TestStats
+from buildscripts.util.teststats import HistoricTaskData, TestRuntime
 from buildscripts.util.taskname import name_generated_task
 from buildscripts.patch_builds.task_generation import (resmoke_commands, TimeoutInfo,
                                                        validate_task_generation_limit)
@@ -436,11 +436,12 @@ def _set_resmoke_cmd(repeat_config: RepeatConfig, resmoke_args: [str]) -> [str]:
     return new_args
 
 
-def _parse_avg_test_runtime(test: str, task_avg_test_runtime_stats: [TestStats]) -> Optional[float]:
+def _parse_avg_test_runtime(test: str,
+                            task_avg_test_runtime_stats: List[TestRuntime]) -> Optional[float]:
     """
-    Parse list of teststats to find runtime for particular test.
+    Parse list of test runtimes to find runtime for particular test.
 
-    :param task_avg_test_runtime_stats: Teststat data.
+    :param task_avg_test_runtime_stats: List of average historic runtimes of tests.
     :param test: Test name.
     :return: Historical average runtime of the test.
     """
@@ -486,13 +487,13 @@ def _calculate_exec_timeout(repeat_config: RepeatConfig, avg_test_runtime: float
 
 
 def _generate_timeouts(repeat_config: RepeatConfig, test: str,
-                       task_avg_test_runtime_stats: [TestStats]) -> TimeoutInfo:
+                       task_avg_test_runtime_stats: [TestRuntime]) -> TimeoutInfo:
     """
     Add timeout.update command to list of commands for a burn in execution task.
 
     :param repeat_config: Information on how the test will repeat.
     :param test: Test name.
-    :param task_avg_test_runtime_stats: Teststat data.
+    :param task_avg_test_runtime_stats: Average historic runtimes of tests.
     :return: TimeoutInfo to use.
     """
     if task_avg_test_runtime_stats:
@@ -512,7 +513,7 @@ def _generate_timeouts(repeat_config: RepeatConfig, test: str,
 
 
 def _get_task_runtime_history(evg_api: Optional[EvergreenApi], project: str, task: str,
-                              variant: str):
+                              variant: str) -> List[TestRuntime]:
     """
     Fetch historical average runtime for all tests in a task from Evergreen API.
 
@@ -528,12 +529,9 @@ def _get_task_runtime_history(evg_api: Optional[EvergreenApi], project: str, tas
     try:
         end_date = datetime.datetime.utcnow().replace(microsecond=0)
         start_date = end_date - datetime.timedelta(days=AVG_TEST_RUNTIME_ANALYSIS_DAYS)
-        data = evg_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=AVG_TEST_RUNTIME_ANALYSIS_DAYS)
-        test_runtimes = TestStats(data).get_tests_runtimes()
-        return test_runtimes
+        test_stats = HistoricTaskData.from_evg(evg_api, project, start_date=start_date,
+                                               end_date=end_date, task=task, variant=variant)
+        return test_stats.get_tests_runtimes()
     except requests.HTTPError as err:
         if err.response.status_code == requests.codes.SERVICE_UNAVAILABLE:
             # Evergreen may return a 503 when the service is degraded.
@@ -544,7 +542,7 @@ def _get_task_runtime_history(evg_api: Optional[EvergreenApi], project: str, tas
 
 
 def _create_task(index: int, test_count: int, test: str, task_data: Dict,
-                 task_runtime_stats: List[TestStats], generate_config: GenerateConfig,
+                 task_runtime_stats: List[TestRuntime], generate_config: GenerateConfig,
                  repeat_config: RepeatConfig, task_prefix: str) -> Task:
     # pylint: disable=too-many-arguments,too-many-locals
     """