# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

import datetime as dt
import urllib
from unittest import mock

import pendulum
import pytest
from sqlalchemy.orm import contains_eager

from airflow.jobs.job import Job
from airflow.jobs.triggerer_job_runner import TriggererJobRunner
from airflow.models import DagRun, SlaMiss, TaskInstance, Trigger
from airflow.models.renderedtifields import RenderedTaskInstanceFields as RTIF
from airflow.security import permissions
from airflow.utils.platform import getuser
from airflow.utils.session import provide_session
from airflow.utils.state import State
from airflow.utils.timezone import datetime
from airflow.utils.types import DagRunType
from tests.test_utils.api_connexion_utils import assert_401, create_user, delete_roles, delete_user
from tests.test_utils.db import clear_db_runs, clear_db_sla_miss, clear_rendered_ti_fields

DEFAULT_DATETIME_1 = datetime(2020, 1, 1)
DEFAULT_DATETIME_STR_1 = "2020-01-01T00:00:00+00:00"
DEFAULT_DATETIME_STR_2 = "2020-01-02T00:00:00+00:00"

QUOTED_DEFAULT_DATETIME_STR_1 = urllib.parse.quote(DEFAULT_DATETIME_STR_1)
QUOTED_DEFAULT_DATETIME_STR_2 = urllib.parse.quote(DEFAULT_DATETIME_STR_2)


@pytest.fixture(scope="module")
def configured_app(minimal_app_for_api):
    app = minimal_app_for_api
    create_user(
        app,  # type: ignore
        username="test",
        role_name="Test",
        permissions=[
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG_RUN),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_TASK_INSTANCE),
            (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_TASK_INSTANCE),
        ],
    )
    create_user(
        app,  # type: ignore
        username="test_dag_read_only",
        role_name="TestDagReadOnly",
        permissions=[
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG_RUN),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_TASK_INSTANCE),
            (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_TASK_INSTANCE),
        ],
    )
    create_user(
        app,  # type: ignore
        username="test_task_read_only",
        role_name="TestTaskReadOnly",
        permissions=[
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_EDIT, permissions.RESOURCE_DAG),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_DAG_RUN),
            (permissions.ACTION_CAN_READ, permissions.RESOURCE_TASK_INSTANCE),
        ],
    )
    create_user(app, username="test_no_permissions", role_name="TestNoPermissions")  # type: ignore

    yield app

    delete_user(app, username="test")  # type: ignore
    delete_user(app, username="test_dag_read_only")  # type: ignore
    delete_user(app, username="test_task_read_only")  # type: ignore
    delete_user(app, username="test_no_permissions")  # type: ignore
    delete_roles(app)


class TestTaskInstanceEndpoint:
    @pytest.fixture(autouse=True)
    def setup_attrs(self, configured_app, dagbag) -> None:
        self.default_time = DEFAULT_DATETIME_1
        self.ti_init = {
            "execution_date": self.default_time,
            "state": State.RUNNING,
        }
        self.ti_extras = {
            "start_date": self.default_time + dt.timedelta(days=1),
            "end_date": self.default_time + dt.timedelta(days=2),
            "pid": 100,
            "duration": 10000,
            "pool": "default_pool",
            "queue": "default_queue",
            "job_id": 0,
        }
        self.app = configured_app
        self.client = self.app.test_client()  # type:ignore
        clear_db_runs()
        clear_db_sla_miss()
        clear_rendered_ti_fields()
        self.dagbag = dagbag

    def create_task_instances(
        self,
        session,
        dag_id: str = "example_python_operator",
        update_extras: bool = True,
        task_instances=None,
        dag_run_state=State.RUNNING,
    ):
        """Method to create task instances using kwargs and default arguments"""

        dag = self.dagbag.get_dag(dag_id)
        tasks = dag.tasks
        counter = len(tasks)
        if task_instances is not None:
            counter = min(len(task_instances), counter)

        run_id = "TEST_DAG_RUN_ID"
        execution_date = self.ti_init.pop("execution_date", self.default_time)
        dr = None

        tis = []
        for i in range(counter):
            if task_instances is None:
                pass
            elif update_extras:
                self.ti_extras.update(task_instances[i])
            else:
                self.ti_init.update(task_instances[i])

            if "execution_date" in self.ti_init:
                run_id = f"TEST_DAG_RUN_ID_{i}"
                execution_date = self.ti_init.pop("execution_date")
                dr = None

            if not dr:
                dr = DagRun(
                    run_id=run_id,
                    dag_id=dag_id,
                    execution_date=execution_date,
                    run_type=DagRunType.MANUAL,
                    state=dag_run_state,
                )
                session.add(dr)
            ti = TaskInstance(task=tasks[i], **self.ti_init)
            ti.dag_run = dr
            ti.note = "placeholder-note"

            for key, value in self.ti_extras.items():
                setattr(ti, key, value)
            session.add(ti)
            tis.append(ti)

        session.commit()
        return tis


class TestGetTaskInstance(TestTaskInstanceEndpoint):
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @pytest.mark.parametrize("username", ["test", "test_dag_read_only", "test_task_read_only"])
    @provide_session
    def test_should_respond_200(self, username, session):
        self.create_task_instances(session)
        # Update ti and set operator to None to
        # test that operator field is nullable.
        # This prevents issue when users upgrade to 2.0+
        # from 1.10.x
        # https://github.com/apache/airflow/issues/14421
        session.query(TaskInstance).update({TaskInstance.operator: None}, synchronize_session="fetch")
        session.commit()
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
            environ_overrides={"REMOTE_USER": username},
        )
        assert response.status_code == 200
        assert response.json == {
            "dag_id": "example_python_operator",
            "duration": 10000.0,
            "end_date": "2020-01-03T00:00:00+00:00",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "executor_config": "{}",
            "hostname": "",
            "map_index": -1,
            "max_tries": 0,
            "note": "placeholder-note",
            "operator": None,
            "pid": 100,
            "pool": "default_pool",
            "pool_slots": 1,
            "priority_weight": 11,
            "queue": "default_queue",
            "queued_when": None,
            "sla_miss": None,
            "start_date": "2020-01-02T00:00:00+00:00",
            "state": "running",
            "task_id": "print_the_context",
            "try_number": 0,
            "unixname": getuser(),
            "dag_run_id": "TEST_DAG_RUN_ID",
            "rendered_fields": {},
            "trigger": None,
            "triggerer_job": None,
        }

    def test_should_respond_200_with_task_state_in_deferred(self, session):
        now = pendulum.now("UTC")
        ti = self.create_task_instances(
            session, task_instances=[{"state": State.DEFERRED}], update_extras=True
        )[0]
        ti.trigger = Trigger("none", {})
        ti.trigger.created_date = now
        ti.triggerer_job = Job()
        TriggererJobRunner(job=ti.triggerer_job)
        ti.triggerer_job.state = "running"
        session.commit()
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
            environ_overrides={"REMOTE_USER": "test"},
        )
        data = response.json

        # this logic in effect replicates mock.ANY for these values
        values_to_ignore = {
            "trigger": ["created_date", "id", "triggerer_id"],
            "triggerer_job": ["executor_class", "hostname", "id", "latest_heartbeat", "start_date"],
        }
        for k, v in values_to_ignore.items():
            for elem in v:
                del data[k][elem]

        assert response.status_code == 200
        assert data == {
            "dag_id": "example_python_operator",
            "duration": 10000.0,
            "end_date": "2020-01-03T00:00:00+00:00",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "executor_config": "{}",
            "hostname": "",
            "map_index": -1,
            "max_tries": 0,
            "note": "placeholder-note",
            "operator": "_PythonDecoratedOperator",
            "pid": 100,
            "pool": "default_pool",
            "pool_slots": 1,
            "priority_weight": 11,
            "queue": "default_queue",
            "queued_when": None,
            "sla_miss": None,
            "start_date": "2020-01-02T00:00:00+00:00",
            "state": "deferred",
            "task_id": "print_the_context",
            "try_number": 0,
            "unixname": getuser(),
            "dag_run_id": "TEST_DAG_RUN_ID",
            "rendered_fields": {},
            "trigger": {
                "classpath": "none",
                "kwargs": "{}",
            },
            "triggerer_job": {
                "dag_id": None,
                "end_date": None,
                "job_type": "TriggererJob",
                "state": "running",
                "unixname": getuser(),
            },
        }

    def test_should_respond_200_with_task_state_in_removed(self, session):
        self.create_task_instances(session, task_instances=[{"state": State.REMOVED}], update_extras=True)
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200
        assert response.json == {
            "dag_id": "example_python_operator",
            "duration": 10000.0,
            "end_date": "2020-01-03T00:00:00+00:00",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "executor_config": "{}",
            "hostname": "",
            "map_index": -1,
            "max_tries": 0,
            "note": "placeholder-note",
            "operator": "_PythonDecoratedOperator",
            "pid": 100,
            "pool": "default_pool",
            "pool_slots": 1,
            "priority_weight": 11,
            "queue": "default_queue",
            "queued_when": None,
            "sla_miss": None,
            "start_date": "2020-01-02T00:00:00+00:00",
            "state": "removed",
            "task_id": "print_the_context",
            "try_number": 0,
            "unixname": getuser(),
            "dag_run_id": "TEST_DAG_RUN_ID",
            "rendered_fields": {},
            "trigger": None,
            "triggerer_job": None,
        }

    def test_should_respond_200_task_instance_with_sla_and_rendered(self, session):
        tis = self.create_task_instances(session)
        session.query()
        sla_miss = SlaMiss(
            task_id="print_the_context",
            dag_id="example_python_operator",
            execution_date=self.default_time,
            timestamp=self.default_time,
        )
        session.add(sla_miss)
        rendered_fields = RTIF(tis[0], render_templates=False)
        session.add(rendered_fields)
        session.commit()
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200

        assert response.json == {
            "dag_id": "example_python_operator",
            "duration": 10000.0,
            "end_date": "2020-01-03T00:00:00+00:00",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "executor_config": "{}",
            "hostname": "",
            "map_index": -1,
            "max_tries": 0,
            "note": "placeholder-note",
            "operator": "_PythonDecoratedOperator",
            "pid": 100,
            "pool": "default_pool",
            "pool_slots": 1,
            "priority_weight": 11,
            "queue": "default_queue",
            "queued_when": None,
            "sla_miss": {
                "dag_id": "example_python_operator",
                "description": None,
                "email_sent": False,
                "execution_date": "2020-01-01T00:00:00+00:00",
                "notification_sent": False,
                "task_id": "print_the_context",
                "timestamp": "2020-01-01T00:00:00+00:00",
            },
            "start_date": "2020-01-02T00:00:00+00:00",
            "state": "running",
            "task_id": "print_the_context",
            "try_number": 0,
            "unixname": getuser(),
            "dag_run_id": "TEST_DAG_RUN_ID",
            "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None},
            "trigger": None,
            "triggerer_job": None,
        }

    def test_should_respond_200_mapped_task_instance_with_rtif(self, session):
        """Verify we don't duplicate rows through join to RTIF"""
        tis = self.create_task_instances(session)
        old_ti = tis[0]
        for idx in (1, 2):
            ti = TaskInstance(task=old_ti.task, run_id=old_ti.run_id, map_index=idx)
            ti.rendered_task_instance_fields = RTIF(ti, render_templates=False)
            for attr in ["duration", "end_date", "pid", "start_date", "state", "queue", "note"]:
                setattr(ti, attr, getattr(old_ti, attr))
            session.add(ti)
        session.commit()

        # in each loop, we should get the right mapped TI back
        for map_index in (1, 2):
            response = self.client.get(
                "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances"
                f"/print_the_context/{map_index}",
                environ_overrides={"REMOTE_USER": "test"},
            )
            assert response.status_code == 200

            assert response.json == {
                "dag_id": "example_python_operator",
                "duration": 10000.0,
                "end_date": "2020-01-03T00:00:00+00:00",
                "execution_date": "2020-01-01T00:00:00+00:00",
                "executor_config": "{}",
                "hostname": "",
                "map_index": map_index,
                "max_tries": 0,
                "note": "placeholder-note",
                "operator": "_PythonDecoratedOperator",
                "pid": 100,
                "pool": "default_pool",
                "pool_slots": 1,
                "priority_weight": 11,
                "queue": "default_queue",
                "queued_when": None,
                "sla_miss": None,
                "start_date": "2020-01-02T00:00:00+00:00",
                "state": "running",
                "task_id": "print_the_context",
                "try_number": 0,
                "unixname": getuser(),
                "dag_run_id": "TEST_DAG_RUN_ID",
                "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None},
                "trigger": None,
                "triggerer_job": None,
            }

    def test_should_raises_401_unauthenticated(self):
        response = self.client.get(
            "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
        )
        assert_401(response)

    def test_should_raise_403_forbidden(self):
        response = self.client.get(
            "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context",
            environ_overrides={"REMOTE_USER": "test_no_permissions"},
        )
        assert response.status_code == 403

    def test_unmapped_map_index_should_return_404(self, session):
        self.create_task_instances(session)
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context/-1",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 404

    def test_should_return_404_for_mapped_endpoint(self, session):
        self.create_task_instances(session)
        for index in ["0", "1", "2"]:
            response = self.client.get(
                "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/"
                f"taskInstances/print_the_context/{index}",
                environ_overrides={"REMOTE_USER": "test"},
            )
            assert response.status_code == 404

    def test_should_return_404_for_list_mapped_endpoint(self, session):
        self.create_task_instances(session)
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/"
            "taskInstances/print_the_context/listMapped",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 404


class TestGetTaskInstances(TestTaskInstanceEndpoint):
    @pytest.mark.parametrize(
        "task_instances, update_extras, url, expected_ti",
        [
            pytest.param(
                [
                    {"execution_date": DEFAULT_DATETIME_1},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                ],
                False,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/~/"
                    f"taskInstances?execution_date_lte={QUOTED_DEFAULT_DATETIME_STR_1}"
                ),
                1,
                id="test execution date filter",
            ),
            pytest.param(
                [
                    {"start_date": DEFAULT_DATETIME_1},
                    {"start_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"start_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                ],
                True,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/~/taskInstances"
                    f"?start_date_gte={QUOTED_DEFAULT_DATETIME_STR_1}&"
                    f"start_date_lte={QUOTED_DEFAULT_DATETIME_STR_2}"
                ),
                2,
                id="test start date filter",
            ),
            pytest.param(
                [
                    {"end_date": DEFAULT_DATETIME_1},
                    {"end_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"end_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                ],
                True,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/~/taskInstances?"
                    f"end_date_gte={QUOTED_DEFAULT_DATETIME_STR_1}&"
                    f"end_date_lte={QUOTED_DEFAULT_DATETIME_STR_2}"
                ),
                2,
                id="test end date filter",
            ),
            pytest.param(
                [
                    {"duration": 100},
                    {"duration": 150},
                    {"duration": 200},
                ],
                True,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/"
                    "taskInstances?duration_gte=100&duration_lte=200"
                ),
                3,
                id="test duration filter",
            ),
            pytest.param(
                [
                    {"duration": 100},
                    {"duration": 150},
                    {"duration": 200},
                ],
                True,
                "/api/v1/dags/~/dagRuns/~/taskInstances?duration_gte=100&duration_lte=200",
                3,
                id="test duration filter ~",
            ),
            pytest.param(
                [
                    {"state": State.RUNNING},
                    {"state": State.QUEUED},
                    {"state": State.SUCCESS},
                    {"state": State.NONE},
                ],
                False,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/"
                    "TEST_DAG_RUN_ID/taskInstances?state=running,queued,none"
                ),
                3,
                id="test state filter",
            ),
            pytest.param(
                [
                    {"state": State.NONE},
                    {"state": State.NONE},
                    {"state": State.NONE},
                    {"state": State.NONE},
                ],
                False,
                ("/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances"),
                4,
                id="test null states with no filter",
            ),
            pytest.param(
                [
                    {"pool": "test_pool_1"},
                    {"pool": "test_pool_2"},
                    {"pool": "test_pool_3"},
                ],
                True,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/"
                    "TEST_DAG_RUN_ID/taskInstances?pool=test_pool_1,test_pool_2"
                ),
                2,
                id="test pool filter",
            ),
            pytest.param(
                [
                    {"pool": "test_pool_1"},
                    {"pool": "test_pool_2"},
                    {"pool": "test_pool_3"},
                ],
                True,
                "/api/v1/dags/~/dagRuns/~/taskInstances?pool=test_pool_1,test_pool_2",
                2,
                id="test pool filter ~",
            ),
            pytest.param(
                [
                    {"queue": "test_queue_1"},
                    {"queue": "test_queue_2"},
                    {"queue": "test_queue_3"},
                ],
                True,
                (
                    "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID"
                    "/taskInstances?queue=test_queue_1,test_queue_2"
                ),
                2,
                id="test queue filter",
            ),
            pytest.param(
                [
                    {"queue": "test_queue_1"},
                    {"queue": "test_queue_2"},
                    {"queue": "test_queue_3"},
                ],
                True,
                "/api/v1/dags/~/dagRuns/~/taskInstances?queue=test_queue_1,test_queue_2",
                2,
                id="test queue filter ~",
            ),
        ],
    )
    def test_should_respond_200(self, task_instances, update_extras, url, expected_ti, session):

        self.create_task_instances(
            session,
            update_extras=update_extras,
            task_instances=task_instances,
        )
        response = self.client.get(url, environ_overrides={"REMOTE_USER": "test"})
        assert response.status_code == 200
        assert response.json["total_entries"] == expected_ti
        assert len(response.json["task_instances"]) == expected_ti

    def test_should_respond_200_for_dag_id_filter(self, session):
        self.create_task_instances(session)
        self.create_task_instances(session, dag_id="example_skip_dag")
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/~/taskInstances",
            environ_overrides={"REMOTE_USER": "test"},
        )

        assert response.status_code == 200
        count = session.query(TaskInstance).filter(TaskInstance.dag_id == "example_python_operator").count()
        assert count == response.json["total_entries"]
        assert count == len(response.json["task_instances"])

    def test_should_raises_401_unauthenticated(self):
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/~/taskInstances",
        )
        assert_401(response)

    def test_should_raise_403_forbidden(self):
        response = self.client.get(
            "/api/v1/dags/example_python_operator/dagRuns/~/taskInstances",
            environ_overrides={"REMOTE_USER": "test_no_permissions"},
        )
        assert response.status_code == 403


class TestGetTaskInstancesBatch(TestTaskInstanceEndpoint):
    @pytest.mark.parametrize(
        "task_instances, update_extras, payload, expected_ti_count, username",
        [
            pytest.param(
                [
                    {"queue": "test_queue_1"},
                    {"queue": "test_queue_2"},
                    {"queue": "test_queue_3"},
                ],
                True,
                {"queue": ["test_queue_1", "test_queue_2"]},
                2,
                "test",
                id="test queue filter",
            ),
            pytest.param(
                [
                    {"pool": "test_pool_1"},
                    {"pool": "test_pool_2"},
                    {"pool": "test_pool_3"},
                ],
                True,
                {"pool": ["test_pool_1", "test_pool_2"]},
                2,
                "test_dag_read_only",
                id="test pool filter",
            ),
            pytest.param(
                [
                    {"state": State.RUNNING},
                    {"state": State.QUEUED},
                    {"state": State.SUCCESS},
                    {"state": State.NONE},
                ],
                False,
                {"state": ["running", "queued", "none"]},
                3,
                "test_task_read_only",
                id="test state filter",
            ),
            pytest.param(
                [
                    {"state": State.NONE},
                    {"state": State.NONE},
                    {"state": State.NONE},
                    {"state": State.NONE},
                ],
                False,
                {},
                4,
                "test_task_read_only",
                id="test dag with null states",
            ),
            pytest.param(
                [
                    {"duration": 100},
                    {"duration": 150},
                    {"duration": 200},
                ],
                True,
                {"duration_gte": 100, "duration_lte": 200},
                3,
                "test",
                id="test duration filter",
            ),
            pytest.param(
                [
                    {"end_date": DEFAULT_DATETIME_1},
                    {"end_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"end_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                ],
                True,
                {
                    "end_date_gte": DEFAULT_DATETIME_STR_1,
                    "end_date_lte": DEFAULT_DATETIME_STR_2,
                },
                2,
                "test_task_read_only",
                id="test end date filter",
            ),
            pytest.param(
                [
                    {"start_date": DEFAULT_DATETIME_1},
                    {"start_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"start_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                ],
                True,
                {
                    "start_date_gte": DEFAULT_DATETIME_STR_1,
                    "start_date_lte": DEFAULT_DATETIME_STR_2,
                },
                2,
                "test_dag_read_only",
                id="test start date filter",
            ),
            pytest.param(
                [
                    {"execution_date": DEFAULT_DATETIME_1},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=4)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=5)},
                ],
                False,
                {
                    "execution_date_gte": DEFAULT_DATETIME_1.isoformat(),
                    "execution_date_lte": (DEFAULT_DATETIME_1 + dt.timedelta(days=2)).isoformat(),
                },
                3,
                "test",
                id="with execution date filter",
            ),
            pytest.param(
                [
                    {"execution_date": DEFAULT_DATETIME_1},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3)},
                ],
                False,
                {
                    "dag_run_ids": ["TEST_DAG_RUN_ID_0", "TEST_DAG_RUN_ID_1"],
                },
                2,
                "test",
                id="test dag run id filter",
            ),
            pytest.param(
                [
                    {"execution_date": DEFAULT_DATETIME_1},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2)},
                    {"execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3)},
                ],
                False,
                {
                    "task_ids": ["print_the_context", "log_sql_query"],
                },
                2,
                "test",
                id="test task id filter",
            ),
        ],
    )
    def test_should_respond_200(
        self, task_instances, update_extras, payload, expected_ti_count, username, session
    ):
        self.create_task_instances(
            session,
            update_extras=update_extras,
            task_instances=task_instances,
        )
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": username},
            json=payload,
        )
        assert response.status_code == 200, response.json
        assert expected_ti_count == response.json["total_entries"]
        assert expected_ti_count == len(response.json["task_instances"])

    @pytest.mark.parametrize(
        "task_instances, payload, expected_ti_count",
        [
            pytest.param(
                [
                    {"task": "test_1"},
                    {"task": "test_2"},
                ],
                {"dag_ids": ["latest_only"]},
                2,
                id="task_instance properties",
            ),
        ],
    )
    @provide_session
    def test_should_respond_200_when_task_instance_properties_are_none(
        self, task_instances, payload, expected_ti_count, session
    ):
        self.ti_extras.update(
            {
                "start_date": None,
                "end_date": None,
                "state": None,
            }
        )
        self.create_task_instances(
            session,
            dag_id="latest_only",
            task_instances=task_instances,
        )
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 200, response.json
        assert expected_ti_count == response.json["total_entries"]
        assert expected_ti_count == len(response.json["task_instances"])

    @pytest.mark.parametrize(
        "payload, expected_ti, total_ti",
        [
            pytest.param(
                {"dag_ids": ["example_python_operator", "example_skip_dag"]},
                19,
                19,
                id="with dag filter",
            ),
        ],
    )
    @provide_session
    def test_should_respond_200_dag_ids_filter(self, payload, expected_ti, total_ti, session):
        self.create_task_instances(session)
        self.create_task_instances(session, dag_id="example_skip_dag")
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 200
        assert len(response.json["task_instances"]) == expected_ti
        assert response.json["total_entries"] == total_ti

    def test_should_raises_401_unauthenticated(self):
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            json={"dag_ids": ["example_python_operator", "example_skip_dag"]},
        )
        assert_401(response)

    def test_should_raise_403_forbidden(self):
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": "test_no_permissions"},
            json={"dag_ids": ["example_python_operator", "example_skip_dag"]},
        )
        assert response.status_code == 403

    def test_should_raise_400_for_no_json(self):
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 400
        assert response.json["detail"] == "Request body must not be empty"

    @pytest.mark.parametrize(
        "payload, expected",
        [
            ({"end_date_lte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
            ({"end_date_gte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
            ({"start_date_lte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
            ({"start_date_gte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
            ({"execution_date_gte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
            ({"execution_date_lte": "2020-11-10T12:42:39.442973"}, "is not a 'date-time'"),
        ],
    )
    @provide_session
    def test_should_raise_400_for_naive_and_bad_datetime(self, payload, expected, session):
        self.create_task_instances(session)
        response = self.client.post(
            "/api/v1/dags/~/dagRuns/~/taskInstances/list",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 400
        assert expected in response.json["detail"]


class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
    @pytest.mark.parametrize(
        "main_dag, task_instances, request_dag, payload, expected_ti",
        [
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                ],
                "example_python_operator",
                {
                    "dry_run": True,
                    "start_date": DEFAULT_DATETIME_STR_2,
                    "only_failed": True,
                },
                2,
                id="clear start date filter",
            ),
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                ],
                "example_python_operator",
                {
                    "dry_run": True,
                    "end_date": DEFAULT_DATETIME_STR_2,
                    "only_failed": True,
                },
                2,
                id="clear end date filter",
            ),
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.RUNNING},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.RUNNING,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                ],
                "example_python_operator",
                {"dry_run": True, "only_running": True, "only_failed": False},
                2,
                id="clear only running",
            ),
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.RUNNING,
                    },
                ],
                "example_python_operator",
                {
                    "dry_run": True,
                    "only_failed": True,
                },
                2,
                id="clear only failed",
            ),
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3),
                        "state": State.FAILED,
                    },
                ],
                "example_python_operator",
                {
                    "dry_run": True,
                    "task_ids": ["print_the_context", "sleep_for_1"],
                },
                2,
                id="clear by task ids",
            ),
            pytest.param(
                "example_subdag_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=4),
                        "state": State.FAILED,
                    },
                ],
                "example_subdag_operator.section-1",
                {"dry_run": True, "include_parentdag": True},
                4,
                id="include parent dag",
            ),
            pytest.param(
                "example_subdag_operator.section-1",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3),
                        "state": State.FAILED,
                    },
                ],
                "example_subdag_operator",
                {
                    "dry_run": True,
                    "include_subdags": True,
                },
                4,
                id="include sub dag",
            ),
            pytest.param(
                "example_python_operator",
                [
                    {"execution_date": DEFAULT_DATETIME_1, "state": State.FAILED},
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                        "state": State.FAILED,
                    },
                    {
                        "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                        "state": State.RUNNING,
                    },
                ],
                "example_python_operator",
                {
                    "only_failed": True,
                },
                2,
                id="dry_run default",
            ),
        ],
    )
    @provide_session
    def test_should_respond_200(self, main_dag, task_instances, request_dag, payload, expected_ti, session):
        self.create_task_instances(
            session,
            dag_id=main_dag,
            task_instances=task_instances,
            update_extras=False,
        )
        self.app.dag_bag.sync_to_db()
        response = self.client.post(
            f"/api/v1/dags/{request_dag}/clearTaskInstances",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 200
        assert len(response.json["task_instances"]) == expected_ti

    @mock.patch("airflow.api_connexion.endpoints.task_instance_endpoint.clear_task_instances")
    def test_clear_taskinstance_is_called_with_queued_dr_state(self, mock_clearti, session):
        """Test that if reset_dag_runs is True, then clear_task_instances is called with State.QUEUED"""
        self.create_task_instances(session)
        dag_id = "example_python_operator"
        payload = {"include_subdags": True, "reset_dag_runs": True, "dry_run": False}
        self.app.dag_bag.sync_to_db()
        response = self.client.post(
            f"/api/v1/dags/{dag_id}/clearTaskInstances",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 200
        mock_clearti.assert_called_once_with(
            [], session, dag=self.app.dag_bag.get_dag(dag_id), dag_run_state=State.QUEUED
        )

    def test_clear_taskinstance_is_called_with_invalid_task_ids(self, session):
        """Test that dagrun is running when invalid task_ids are passed to clearTaskInstances API."""
        dag_id = "example_python_operator"
        tis = self.create_task_instances(session)
        dagrun = tis[0].get_dagrun()
        assert dagrun.state == "running"

        payload = {"dry_run": False, "reset_dag_runs": True, "task_ids": [""]}
        self.app.dag_bag.sync_to_db()
        response = self.client.post(
            f"/api/v1/dags/{dag_id}/clearTaskInstances",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 200

        dagrun.refresh_from_db()
        assert dagrun.state == "running"
        assert all(ti.state == "running" for ti in tis)

    def test_should_respond_200_with_reset_dag_run(self, session):
        dag_id = "example_python_operator"
        payload = {
            "dry_run": False,
            "reset_dag_runs": True,
            "only_failed": False,
            "only_running": True,
            "include_subdags": True,
        }
        task_instances = [
            {"execution_date": DEFAULT_DATETIME_1, "state": State.RUNNING},
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                "state": State.RUNNING,
            },
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=2),
                "state": State.RUNNING,
            },
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=3),
                "state": State.RUNNING,
            },
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=4),
                "state": State.RUNNING,
            },
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=5),
                "state": State.RUNNING,
            },
        ]

        self.create_task_instances(
            session,
            dag_id=dag_id,
            task_instances=task_instances,
            update_extras=False,
            dag_run_state=State.FAILED,
        )
        response = self.client.post(
            f"/api/v1/dags/{dag_id}/clearTaskInstances",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )

        failed_dag_runs = session.query(DagRun).filter(DagRun.state == "failed").count()
        assert 200 == response.status_code
        expected_response = [
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_0",
                "execution_date": "2020-01-01T00:00:00+00:00",
                "task_id": "print_the_context",
            },
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_1",
                "execution_date": "2020-01-02T00:00:00+00:00",
                "task_id": "log_sql_query",
            },
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_2",
                "execution_date": "2020-01-03T00:00:00+00:00",
                "task_id": "sleep_for_0",
            },
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_3",
                "execution_date": "2020-01-04T00:00:00+00:00",
                "task_id": "sleep_for_1",
            },
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_4",
                "execution_date": "2020-01-05T00:00:00+00:00",
                "task_id": "sleep_for_2",
            },
            {
                "dag_id": "example_python_operator",
                "dag_run_id": "TEST_DAG_RUN_ID_5",
                "execution_date": "2020-01-06T00:00:00+00:00",
                "task_id": "sleep_for_3",
            },
        ]
        for task_instance in expected_response:
            assert task_instance in response.json["task_instances"]
        assert 6 == len(response.json["task_instances"])
        assert 0 == failed_dag_runs, 0

    def test_should_raises_401_unauthenticated(self):
        response = self.client.post(
            "/api/v1/dags/example_python_operator/clearTaskInstances",
            json={
                "dry_run": False,
                "reset_dag_runs": True,
                "only_failed": False,
                "only_running": True,
                "include_subdags": True,
            },
        )
        assert_401(response)

    @pytest.mark.parametrize("username", ["test_no_permissions", "test_dag_read_only", "test_task_read_only"])
    def test_should_raise_403_forbidden(self, username: str):
        response = self.client.post(
            "/api/v1/dags/example_python_operator/clearTaskInstances",
            environ_overrides={"REMOTE_USER": username},
            json={
                "dry_run": False,
                "reset_dag_runs": True,
                "only_failed": False,
                "only_running": True,
                "include_subdags": True,
            },
        )
        assert response.status_code == 403

    @pytest.mark.parametrize(
        "payload, expected",
        [
            ({"end_date": "2020-11-10T12:42:39.442973"}, "Naive datetime is disallowed"),
            ({"end_date": "2020-11-10T12:4po"}, "{'end_date': ['Not a valid datetime.']}"),
            ({"start_date": "2020-11-10T12:42:39.442973"}, "Naive datetime is disallowed"),
            ({"start_date": "2020-11-10T12:4po"}, "{'start_date': ['Not a valid datetime.']}"),
        ],
    )
    @provide_session
    def test_should_raise_400_for_naive_and_bad_datetime(self, payload, expected, session):
        task_instances = [
            {"execution_date": DEFAULT_DATETIME_1, "state": State.RUNNING},
            {
                "execution_date": DEFAULT_DATETIME_1 + dt.timedelta(days=1),
                "state": State.RUNNING,
            },
        ]
        self.create_task_instances(
            session,
            dag_id="example_python_operator",
            task_instances=task_instances,
            update_extras=False,
        )
        self.app.dag_bag.sync_to_db()
        response = self.client.post(
            "/api/v1/dags/example_python_operator/clearTaskInstances",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 400
        assert response.json["detail"] == expected


class TestPostSetTaskInstanceState(TestTaskInstanceEndpoint):
    @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
    def test_should_assert_call_mocked_api(self, mock_set_task_instance_state, session):
        self.create_task_instances(session)
        mock_set_task_instance_state.return_value = (
            session.query(TaskInstance)
            .join(TaskInstance.dag_run)
            .options(contains_eager(TaskInstance.dag_run))
            .filter(TaskInstance.task_id == "print_the_context")
            .all()
        )
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "execution_date": DEFAULT_DATETIME_1.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 200
        assert response.json == {
            "task_instances": [
                {
                    "dag_id": "example_python_operator",
                    "dag_run_id": "TEST_DAG_RUN_ID",
                    "execution_date": "2020-01-01T00:00:00+00:00",
                    "task_id": "print_the_context",
                }
            ]
        }

        mock_set_task_instance_state.assert_called_once_with(
            commit=False,
            downstream=True,
            run_id=None,
            execution_date=DEFAULT_DATETIME_1,
            future=True,
            past=True,
            state="failed",
            task_id="print_the_context",
            upstream=True,
            session=session,
        )

    @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
    def test_should_assert_call_mocked_api_when_run_id(self, mock_set_task_instance_state, session):
        self.create_task_instances(session)
        run_id = "TEST_DAG_RUN_ID"
        mock_set_task_instance_state.return_value = (
            session.query(TaskInstance)
            .join(TaskInstance.dag_run)
            .filter(TaskInstance.task_id == "print_the_context")
            .all()
        )
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "dag_run_id": run_id,
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 200
        assert response.json == {
            "task_instances": [
                {
                    "dag_id": "example_python_operator",
                    "dag_run_id": "TEST_DAG_RUN_ID",
                    "execution_date": "2020-01-01T00:00:00+00:00",
                    "task_id": "print_the_context",
                }
            ]
        }

        mock_set_task_instance_state.assert_called_once_with(
            commit=False,
            downstream=True,
            run_id=run_id,
            execution_date=None,
            future=True,
            past=True,
            state="failed",
            task_id="print_the_context",
            upstream=True,
            session=session,
        )

    @pytest.mark.parametrize(
        "error, code, payload",
        [
            [
                "{'_schema': ['Exactly one of execution_date or dag_run_id must be provided']}",
                400,
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
            ],
            [
                "Task instance not found for task 'print_the_context' on execution_date "
                "2021-01-01 00:00:00+00:00",
                404,
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "execution_date": "2021-01-01T00:00:00+00:00",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
            ],
            [
                "Task instance not found for task 'print_the_context' on DAG run with ID 'TEST_DAG_RUN_'",
                404,
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "dag_run_id": "TEST_DAG_RUN_",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
            ],
            [
                "{'_schema': ['Exactly one of execution_date or dag_run_id must be provided']}",
                400,
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "dag_run_id": "TEST_DAG_RUN_",
                    "execution_date": "2020-01-01T00:00:00+00:00",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
            ],
        ],
    )
    def test_should_handle_errors(self, error, code, payload, session):
        self.create_task_instances(session)
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == code
        assert response.json["detail"] == error

    def test_should_raises_401_unauthenticated(self):
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "execution_date": DEFAULT_DATETIME_1.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert_401(response)

    @pytest.mark.parametrize("username", ["test_no_permissions", "test_dag_read_only", "test_task_read_only"])
    def test_should_raise_403_forbidden(self, username):
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": username},
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "execution_date": DEFAULT_DATETIME_1.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 403

    def test_should_raise_404_not_found_dag(self):
        response = self.client.post(
            "/api/v1/dags/INVALID_DAG/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "execution_date": DEFAULT_DATETIME_1.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 404

    @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
    def test_should_raise_not_found_if_execution_date_is_wrong(self, mock_set_task_instance_state, session):
        self.create_task_instances(session)
        date = DEFAULT_DATETIME_1 + dt.timedelta(days=1)
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "task_id": "print_the_context",
                "execution_date": date.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 404
        assert response.json["detail"] == (
            f"Task instance not found for task 'print_the_context' on execution_date {date}"
        )
        assert mock_set_task_instance_state.call_count == 0

    def test_should_raise_404_not_found_task(self):
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "task_id": "INVALID_TASK",
                "execution_date": DEFAULT_DATETIME_1.isoformat(),
                "include_upstream": True,
                "include_downstream": True,
                "include_future": True,
                "include_past": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 404

    @pytest.mark.parametrize(
        "payload, expected",
        [
            (
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "execution_date": "2020-11-10T12:42:39.442973",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
                "Naive datetime is disallowed",
            ),
            (
                {
                    "dry_run": True,
                    "task_id": "print_the_context",
                    "execution_date": "2020-11-10T12:4opfo",
                    "include_upstream": True,
                    "include_downstream": True,
                    "include_future": True,
                    "include_past": True,
                    "new_state": "failed",
                },
                "{'execution_date': ['Not a valid datetime.']}",
            ),
        ],
    )
    @provide_session
    def test_should_raise_400_for_naive_and_bad_datetime(self, payload, expected, session):
        self.create_task_instances(session)
        response = self.client.post(
            "/api/v1/dags/example_python_operator/updateTaskInstancesState",
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 400
        assert response.json["detail"] == expected


class TestPatchTaskInstance(TestTaskInstanceEndpoint):
    ENDPOINT_URL = (
        "/api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
    )

    @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
    def test_should_call_mocked_api(self, mock_set_task_instance_state, session):
        self.create_task_instances(session)

        NEW_STATE = "failed"
        mock_set_task_instance_state.return_value = session.get(
            TaskInstance,
            {
                "task_id": "print_the_context",
                "dag_id": "example_python_operator",
                "run_id": "TEST_DAG_RUN_ID",
                "map_index": -1,
            },
        )
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": False,
                "new_state": NEW_STATE,
            },
        )
        assert response.status_code == 200
        assert response.json == {
            "dag_id": "example_python_operator",
            "dag_run_id": "TEST_DAG_RUN_ID",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "task_id": "print_the_context",
        }

        mock_set_task_instance_state.assert_called_once_with(
            task_id="print_the_context",
            run_id="TEST_DAG_RUN_ID",
            map_indexes=[-1],
            state=NEW_STATE,
            commit=True,
            session=session,
        )

    @mock.patch("airflow.models.dag.DAG.set_task_instance_state")
    def test_should_not_call_mocked_api_for_dry_run(self, mock_set_task_instance_state, session):
        self.create_task_instances(session)

        NEW_STATE = "failed"
        mock_set_task_instance_state.return_value = session.get(
            TaskInstance,
            {
                "task_id": "print_the_context",
                "dag_id": "example_python_operator",
                "run_id": "TEST_DAG_RUN_ID",
                "map_index": -1,
            },
        )
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "new_state": NEW_STATE,
            },
        )
        assert response.status_code == 200
        print(response.status_code)
        assert response.json == {
            "dag_id": "example_python_operator",
            "dag_run_id": "TEST_DAG_RUN_ID",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "task_id": "print_the_context",
        }

        mock_set_task_instance_state.assert_not_called()

    def test_should_update_task_instance_state(self, session):
        self.create_task_instances(session)

        NEW_STATE = "failed"

        self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": False,
                "new_state": NEW_STATE,
            },
        )

        response2 = self.client.get(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={},
        )
        assert response2.status_code == 200
        assert response2.json["state"] == NEW_STATE

    def test_should_update_mapped_task_instance_state(self, session):
        NEW_STATE = "failed"
        map_index = 1
        tis = self.create_task_instances(session)
        ti = TaskInstance(task=tis[0].task, run_id=tis[0].run_id, map_index=map_index)
        ti.rendered_task_instance_fields = RTIF(ti, render_templates=False)
        session.add(ti)
        session.commit()

        self.client.patch(
            f"{self.ENDPOINT_URL}/{map_index}",
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": False,
                "new_state": NEW_STATE,
            },
        )

        response2 = self.client.get(
            f"{self.ENDPOINT_URL}/{map_index}",
            environ_overrides={"REMOTE_USER": "test"},
            json={},
        )
        assert response2.status_code == 200
        assert response2.json["state"] == NEW_STATE

    @pytest.mark.parametrize(
        "error, code, payload",
        [
            [
                "Task instance not found for task 'print_the_context' on DAG run with ID 'TEST_DAG_RUN_ID'",
                404,
                {
                    "dry_run": True,
                    "new_state": "failed",
                },
            ]
        ],
    )
    def test_should_handle_errors(self, error, code, payload, session):
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == code
        assert response.json["detail"] == error

    def test_should_raises_401_unauthenticated(self):
        response = self.client.patch(
            self.ENDPOINT_URL,
            json={
                "dry_run": False,
                "new_state": "failed",
            },
        )
        assert_401(response)

    @pytest.mark.parametrize("username", ["test_no_permissions", "test_dag_read_only", "test_task_read_only"])
    def test_should_raise_403_forbidden(self, username):
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": username},
            json={
                "dry_run": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 403

    def test_should_raise_404_not_found_dag(self):
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 404

    def test_should_raise_404_not_found_task(self):
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json={
                "dry_run": True,
                "new_state": "failed",
            },
        )
        assert response.status_code == 404

    @pytest.mark.parametrize(
        "payload, expected",
        [
            (
                {
                    "dry_run": True,
                    "new_state": "failede",
                },
                f"'failede' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']"
                " - 'new_state'",
            ),
            (
                {
                    "dry_run": True,
                    "new_state": "queued",
                },
                f"'queued' is not one of ['{State.SUCCESS}', '{State.FAILED}', '{State.SKIPPED}']"
                " - 'new_state'",
            ),
        ],
    )
    @provide_session
    def test_should_raise_400_for_invalid_task_instance_state(self, payload, expected, session):
        self.create_task_instances(session)
        response = self.client.patch(
            self.ENDPOINT_URL,
            environ_overrides={"REMOTE_USER": "test"},
            json=payload,
        )
        assert response.status_code == 400
        assert response.json["detail"] == expected
        assert response.json["detail"] == expected


class TestSetTaskInstanceNote(TestTaskInstanceEndpoint):
    def setup_method(self):
        clear_db_runs()

    def teardown_method(self):
        clear_db_runs()

    @provide_session
    def test_should_respond_200(self, session):
        tis = self.create_task_instances(session)
        new_note_value = "My super cool TaskInstance note."
        response = self.client.patch(
            "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/"
            "print_the_context/setNote",
            json={"note": new_note_value},
            environ_overrides={"REMOTE_USER": "test"},
        )
        assert response.status_code == 200, response.text
        assert response.json == {
            "dag_id": "example_python_operator",
            "duration": 10000.0,
            "end_date": "2020-01-03T00:00:00+00:00",
            "execution_date": "2020-01-01T00:00:00+00:00",
            "executor_config": "{}",
            "hostname": "",
            "map_index": -1,
            "max_tries": 0,
            "note": new_note_value,
            "operator": "_PythonDecoratedOperator",
            "pid": 100,
            "pool": "default_pool",
            "pool_slots": 1,
            "priority_weight": 11,
            "queue": "default_queue",
            "queued_when": None,
            "sla_miss": None,
            "start_date": "2020-01-02T00:00:00+00:00",
            "state": "running",
            "task_id": "print_the_context",
            "try_number": 0,
            "unixname": getuser(),
            "dag_run_id": "TEST_DAG_RUN_ID",
            "rendered_fields": {},
            "trigger": None,
            "triggerer_job": None,
        }
        ti = tis[0]
        assert ti.task_instance_note.user_id is not None

    def test_should_respond_200_mapped_task_instance_with_rtif(self, session):
        """Verify we don't duplicate rows through join to RTIF"""
        tis = self.create_task_instances(session)
        old_ti = tis[0]
        for idx in (1, 2):
            ti = TaskInstance(task=old_ti.task, run_id=old_ti.run_id, map_index=idx)
            ti.rendered_task_instance_fields = RTIF(ti, render_templates=False)
            for attr in ["duration", "end_date", "pid", "start_date", "state", "queue", "note"]:
                setattr(ti, attr, getattr(old_ti, attr))
            session.add(ti)
        session.commit()

        # in each loop, we should get the right mapped TI back
        for map_index in (1, 2):
            new_note_value = f"My super cool TaskInstance note {map_index}"
            response = self.client.patch(
                "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/"
                f"print_the_context/{map_index}/setNote",
                json={"note": new_note_value},
                environ_overrides={"REMOTE_USER": "test"},
            )
            assert response.status_code == 200, response.text

            assert response.json == {
                "dag_id": "example_python_operator",
                "duration": 10000.0,
                "end_date": "2020-01-03T00:00:00+00:00",
                "execution_date": "2020-01-01T00:00:00+00:00",
                "executor_config": "{}",
                "hostname": "",
                "map_index": map_index,
                "max_tries": 0,
                "note": new_note_value,
                "operator": "_PythonDecoratedOperator",
                "pid": 100,
                "pool": "default_pool",
                "pool_slots": 1,
                "priority_weight": 11,
                "queue": "default_queue",
                "queued_when": None,
                "sla_miss": None,
                "start_date": "2020-01-02T00:00:00+00:00",
                "state": "running",
                "task_id": "print_the_context",
                "try_number": 0,
                "unixname": getuser(),
                "dag_run_id": "TEST_DAG_RUN_ID",
                "rendered_fields": {"op_args": [], "op_kwargs": {}, "templates_dict": None},
                "trigger": None,
                "triggerer_job": None,
            }

    def test_should_raises_401_unauthenticated(self):
        for map_index in ["", "/0"]:
            url = (
                "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/"
                f"print_the_context{map_index}/setNote"
            )
            response = self.client.patch(
                url,
                json={"note": "I am setting a note while being unauthenticated."},
            )
            assert_401(response)

    def test_should_raise_403_forbidden(self):
        for map_index in ["", "/0"]:
            response = self.client.patch(
                "api/v1/dags/example_python_operator/dagRuns/TEST_DAG_RUN_ID/taskInstances/"
                f"print_the_context{map_index}/setNote",
                json={"note": "I am setting a note without the proper permissions."},
                environ_overrides={"REMOTE_USER": "test_no_permissions"},
            )
            assert response.status_code == 403

    def test_should_respond_404(self, session):
        self.create_task_instances(session)
        for map_index in ["", "/0"]:
            response = self.client.patch(
                f"api/v1/dags/INVALID_DAG_ID/dagRuns/TEST_DAG_RUN_ID/taskInstances/print_the_context"
                f"{map_index}/setNote",
                json={"note": "I am setting a note on a DAG that doesn't exist."},
                environ_overrides={"REMOTE_USER": "test"},
            )
            assert response.status_code == 404
