"""
    Airflow API (Stable)

    Apache Airflow management API.  # noqa: E501

    The version of the OpenAPI document: 1.0.0
    Contact: zach.z.liu@gmail.com
    Generated by: https://openapi-generator.tech
"""


from test.integration.conftest import BCOLORS

import pytest

from airflow_python_sdk.model.update_task_instances_state import \
    UpdateTaskInstancesState

@pytest.mark.parametrize(
    "test_input, expected",
    [
        ([False, False, False, True], None),
        ([True, False, False, True], None),
        ([True, True, False, True], None),
        ([True, True, True, True], None),
        ([True, False, True, True], None),
    ],
)
def test_set_task_instances_state(dag_api_setup, test_input, expected):
    """Test the /dags/{dag_id}/updateTaskInstancesState API EP
    """
    dag_id = "test_glue_partitions_sensor"
    execution_date = "2020-03-25T00:00:00+00:00"
    include_downstream, include_future, include_past, include_upstream = \
        test_input
    update_task_instances_state = UpdateTaskInstancesState(
        dry_run=True,
        execution_date=execution_date,
        include_downstream=include_downstream,
        include_future=include_future,
        include_past=include_past,
        include_upstream=include_upstream,
        new_state="success",
        task_id="task1",
    )  # UpdateTaskInstancesState | Parameters of action

    # Set a state of task instances
    api_response = dag_api_setup.post_set_task_instances_state(
        dag_id, update_task_instances_state
    )
    print(f"{BCOLORS.OKGREEN}OK{BCOLORS.ENDC}")
