mongo/buildscripts/patch_builds/selected_tests/selected_tests_client.py

#!/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)