Source code for operators.start_reflect_run_test

"""
Operator to validate and start a reflect run test.

TODO: add more customizability to the operator (params, etc.)
TODO: complete
TODO: validate (when our monthly limits are not close to nil)
"""

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from airflow.hooks.base import BaseHook

from config import Config




class StartReflectRunOperator(BaseOperator):
    """
    Custom operator to start a reflect run in the specified environment. Uses `requests`.

    :param test_id: The ID of the test to run.
    :type test_id: str
    """

    ui_color = "#FFFFFF"

    @apply_defaults
    def __init__(
        self: object,
        test_id: str,
        base_url: str = Config.REFLECT_RUN_BASE_URL,
        *args,
        **kwargs,
    ):
        super().__init__(
            *args,
            **kwargs,
        )
        self.test_id = test_id
        self.base_url = base_url



    def validate_test_id(
        self: object,
        test_id: str,
    ) -> bool:
        """
        Validate the test ID format.

        :param test_id: The ID of the test to run.
        :type test_id: str

        :return: True if valid, False otherwise.
        :rtype: bool
        """
        import requests as re

        endpoint_url = f"{self.base_url}/tests"
        test_id = int(self.test_id)  # Ensure test_id is an integer

        conn = BaseHook.get_connection(Config.REFLECT_RUN_API_KEY)
        reflect_api_key = conn.extra

        params = {
            "x-api-key": f"{reflect_api_key}",
        }

        tests = self.send_request(
            method="GET",
            endpoint=endpoint_url,
            params=params,
        ).json()

        for test in tests:
            if test["id"] == test_id:
                return True
        return False




    def start_test(
        self: object,
        test_id: str,
    ):
        """
        Start a reflect test run.

        :param test_id: The ID of the test to run.
        :type test_id: str
        """
        import requests as re

        endpoint_url = f"{self.base_url}/tests/{test_id}/executions"
        test_id = int(self.test_id)

        conn = BaseHook.get_connection(Config.REFLECT_RUN_API_KEY)
        reflect_api_key = conn.extra

        params = {
            "x-api-key": f"{reflect_api_key}",
        }

        response = self.send_request(
            method="POST",
            endpoint=endpoint_url,
            params=params,
        )

        print("Execution ID: ", response.json()["executionId"])
        if response.status_code != 200:
            raise Exception(
                f"Failed to start reflect run: {response.status_code} - {response.text}"
            )
        return response.json()["executionId"]




    def execute(
        self: object,
        test_id: str,
        context: Context,
    ):
        """
        Execute the operator.

        :param test_id: The ID of the test to run.
        :type test_id: str
        :param context: The context of the task instance.
        :type context: Context
        """
        if not self.validate_test_id(test_id):
            raise ValueError(f"Invalid test ID: {test_id}")

        execution_id = self.start_test(test_id)
        return execution_id