1 files changed, 171 insertions, 0 deletions
diff --git a/buildscripts/patch_builds/selected_tests/selected_tests_client.py b/buildscripts/patch_builds/selected_tests/selected_tests_client.py
new file mode 100644
index 00000000000..afb50730df4
--- /dev/null
+++ b/buildscripts/patch_builds/selected_tests/selected_tests_client.py
@@ -0,0 +1,171 @@
+#!/usr/bin/env python3
+"""Client for accessing selected test app."""
+
+from typing import Set, List, Dict, Any
+from urllib.parse import urlparse
+
+import requests
+import yaml
+from pydantic import BaseModel
+
+
+class TestFileInstance(BaseModel):
+    """
+    Frequency of how often a test file was seen.
+
+    name: Name of test file.
+    test_file_seen_count: Occurrences of test file.
+    """
+
+    name: str
+    test_file_seen_count: int
+
+
+class TestMapping(BaseModel):
+    """
+    How tests map to the specified source files.
+
+    branch: Git branch being analyzed.
+    project: Evergreen project being analyzed.
+    repo: Git repo being analyzed.
+    source_file: Source file of mappings.
+    source_file_seen_count: Number of occurrences of source file.
+    test_files: Test files that have been changed with the source file.
+    """
+
+    branch: str
+    project: str
+    repo: str
+    source_file: str
+    source_file_seen_count: int
+    test_files: List[TestFileInstance]
+
+
+class TestMappingsResponse(BaseModel):
+    """
+    Response from the test mappings end point.
+
+    test_mappings: List of source files with correlated test files.
+    """
+
+    test_mappings: List[TestMapping]
+
+
+class TaskMapInstance(BaseModel):
+    """
+    Frequency of how often a task is impacted by a source file change.
+
+    name: Name of task that was impacted.
+    variant: Name of build variant task was run on.
+    flip_count: Number of times the task was impacted by the source file.
+    """
+
+    name: str
+    variant: str
+    flip_count: int
+
+
+class TaskMapping(BaseModel):
+    """
+    How tasks map to the specified source files.
+
+    branch: Git branch being analyzed.
+    project: Evergreen project being analyzed.
+    repo: Git repo being analyzed.
+    source_file: Source file of mappings.
+    source_file_seen_count: Number of occurrences of source file.
+    tasks: Tasks that have been impacted by the source file.
+    """
+
+    branch: str
+    project: str
+    repo: str
+    source_file: str
+    source_file_seen_count: int
+    tasks: List[TaskMapInstance]
+
+
+class TaskMappingsResponse(BaseModel):
+    """
+    Response from the task mappings end point.
+
+    task_mappings: List of source files with correlated tasks.
+    """
+
+    task_mappings: List[TaskMapping]
+
+
+class SelectedTestsClient(object):
+    """Selected-tests client object."""
+
+    def __init__(self, url: str, project: str, auth_user: str, auth_token: str) -> None:
+        """
+        Create selected-tests client object.
+
+        :param url: Selected-tests service url.
+        :param project: Selected-tests service project.
+        :param auth_user: Selected-tests service auth user to authenticate request.
+        :param auth_token: Selected-tests service auth token to authenticate request.
+        """
+        self.url = url
+        self.project = project
+        self.session = requests.Session()
+        adapter = requests.adapters.HTTPAdapter()
+        self.session.mount(f"{urlparse(self.url).scheme}://", adapter)
+        self.session.cookies.update({"auth_user": auth_user, "auth_token": auth_token})
+        self.session.headers.update(
+            {"Content-type": "application/json", "Accept": "application/json"})
+
+    @classmethod
+    def from_file(cls, filename: str) -> "SelectedTestsClient":
+        """
+        Read config from given filename.
+
+        :param filename: Filename to read config.
+        :return: Config read from file.
+        """
+        with open(filename, 'r') as fstream:
+            config = yaml.safe_load(fstream)
+            if config:
+                return cls(config["url"], config["project"], config["auth_user"],
+                           config["auth_token"])
+
+        raise ValueError(f"Could not determine selected tests configuration from {filename}")
+
+    def _call_api(self, endpoint: str, payload: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Make a call to the selected tests service and return the response.
+
+        :param endpoint: Endpoint to call.
+        :param payload: Payload to call with.
+        :return: Response from service.
+        """
+        url = f"{self.url}{endpoint}"
+        response = self.session.get(url, params=payload)
+        response.raise_for_status()
+
+        return response.json()
+
+    def get_test_mappings(self, threshold: float, changed_files: Set[str]) -> TestMappingsResponse:
+        """
+        Request related test files from selected-tests service.
+
+        :param threshold: Threshold for test file correlation.
+        :param changed_files: Set of changed_files.
+        :return: Related test files returned by selected-tests service.
+        """
+        payload = {"threshold": threshold, "changed_files": ",".join(changed_files)}
+        response = self._call_api(f"/projects/{self.project}/test-mappings", payload)
+        return TestMappingsResponse(**response)
+
+    def get_task_mappings(self, threshold: float, changed_files: Set[str]) -> TaskMappingsResponse:
+        """
+        Request related tasks from selected-tests service.
+
+        :param threshold: Threshold for test file correlation.
+        :param changed_files: Set of changed_files.
+        :return: Related tasks returned by selected-tests service.
+        """
+        payload = {"threshold": threshold, "changed_files": ",".join(changed_files)}
+        response = self._call_api(f"/projects/{self.project}/task-mappings", payload)
+        return TaskMappingsResponse(**response)