text: # 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.
  import datetime as dt
  from unittest import mock
  
  import pytest
  from parameterized import parameterized
  from sqlalchemy.orm import contains_eager
  
  from airflow.models import DagRun, SlaMiss, TaskInstance
  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"
  
  
  @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
  
              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):
      @parameterized.expand(["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,
              "operator": None,
              "pid": 100,
              "pool": "default_pool",
              "pool_slots": 1,
              "priority_weight": 6,
              "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": {},
          }
  
      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,
              "operator": "_PythonDecoratedOperator",
              "pid": 100,
              "pool": "default_pool",
              "pool_slots": 1,
              "priority_weight": 6,
              "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": {},
          }
  
      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,
              "operator": "_PythonDecoratedOperator",
              "pid": 100,
              "pool": "default_pool",
              "pool_slots": 1,
              "priority_weight": 6,
              "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': {}},
          }
  
      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)
          session.query()
          ti = tis[0]
          ti.map_index = 1
          rendered_fields = RTIF(ti, render_templates=False)
          session.add(rendered_fields)
          session.commit()
          new_ti = TaskInstance(task=ti.task, run_id=ti.run_id, map_index=2)
          for attr in ['duration', 'end_date', 'pid', 'start_date', 'state', 'queue']:
              setattr(new_ti, attr, getattr(ti, attr))
          session.add(new_ti)
          rendered_fields = RTIF(new_ti, render_templates=False)
          session.add(rendered_fields)
          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,
                  "operator": "_PythonDecoratedOperator",
                  "pid": 100,
                  "pool": "default_pool",
                  "pool_slots": 1,
                  "priority_weight": 6,
                  "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': {}},
              }
  
      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
  
  
  class TestGetTaskInstances(TestTaskInstanceEndpoint):
      @parameterized.expand(
          [
              (
                  "test execution date filter",
                  [
                      {"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={DEFAULT_DATETIME_STR_1}"
                  ),
                  1,
              ),
              (
                  "test start date filter",
                  [
                      {"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={DEFAULT_DATETIME_STR_1}&start_date_lte={DEFAULT_DATETIME_STR_2}"
                  ),
                  2,
              ),
              (
                  "test end date filter",
                  [
                      {"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={DEFAULT_DATETIME_STR_1}&end_date_lte={DEFAULT_DATETIME_STR_2}"
                  ),
                  2,
              ),
              (
                  "test duration filter",
                  [
                      {"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,
              ),
              (
                  "test duration filter ~",
                  [
                      {"duration": 100},
                      {"duration": 150},
                      {"duration": 200},
                  ],
                  True,
                  "/api/v1/dags/~/dagRuns/~/taskInstances?duration_gte=100&duration_lte=200",
                  3,
              ),
              (
                  "test state filter",
                  [
                      {"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,
              ),
              (
                  "test null states with no filter",
                  [
                      {"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,
              ),
              (
                  "test pool filter",
                  [
                      {"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,
              ),
              (
                  "test pool filter ~",
                  [
                      {"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,
              ),
              (
                  "test queue filter",
                  [
                      {"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,
              ),
              (
                  "test queue filter ~",
                  [
                      {"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,
              ),
          ]
      )
      @provide_session
      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):
      @parameterized.expand(
          [
              (
                  "test queue filter",
                  [
                      {"queue": "test_queue_1"},
                      {"queue": "test_queue_2"},
                      {"queue": "test_queue_3"},
                  ],
                  True,
                  {"queue": ["test_queue_1", "test_queue_2"]},
                  2,
                  "test",
              ),
              (
                  "test pool filter",
                  [
                      {"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",
              ),
              (
                  "test state filter",
                  [
                      {"state": State.RUNNING},
                      {"state": State.QUEUED},
                      {"state": State.SUCCESS},
                      {"state": State.NONE},
                  ],
                  False,
                  {"state": ["running", "queued", "none"]},
                  3,
                  "test_task_read_only",
              ),
              (
                  "test dag with null states",
                  [
                      {"state": State.NONE},
                      {"state": State.NONE},
                      {"state": State.NONE},
                      {"state": State.NONE},
                  ],
                  False,
                  {},
                  4,
                  "test_task_read_only",
              ),
              (
                  "test duration filter",
                  [
                      {"duration": 100},
                      {"duration": 150},
                      {"duration": 200},
                  ],
                  True,
                  {"duration_gte": 100, "duration_lte": 200},
                  3,
                  "test",
              ),
              (
                  "test end date filter",
                  [
                      {"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",
              ),
              (
                  "test start date filter",
                  [
                      {"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",
              ),
              (
                  "with execution date filter",
                  [
                      {"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,
                      "execution_date_lte": (DEFAULT_DATETIME_1 + dt.timedelta(days=2)),
                  },
                  3,
                  "test",
              ),
          ]
      )
      @provide_session
      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"])
  
      @parameterized.expand(
          [
              (
                  "task_instance properties",
                  [
                      {"task": "test_1"},
                      {"task": "test_2"},
                  ],
                  {"dag_ids": ["latest_only"]},
                  2,
              ),
          ],
      )
      @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"])
  
      @parameterized.expand(
          [
              (
                  "with dag filter",
                  {"dag_ids": ["example_python_operator", "example_skip_dag"]},
                  15,
                  15,
              ),
          ],
      )
      @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
  
      @parameterized.expand(
          [
              ({"end_date_lte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
              ({"end_date_gte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
              ({"start_date_lte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
              ({"start_date_gte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
              ({"execution_date_gte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
              ({"execution_date_lte": '2020-11-10T12:42:39.442973'}, "Naive datetime is disallowed"),
          ]
      )
      @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 response.json['detail'] == expected
  
  
  class TestPostClearTaskInstances(TestTaskInstanceEndpoint):
      @parameterized.expand(
          [
              (
                  "clear start date filter",
                  "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,
              ),
              (
                  "clear end date filter",
                  "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,
              ),
              (
                  "clear only running",
                  "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,
              ),
              (
                  "clear only failed",
                  "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,
              ),
              (
                  "clear by task ids",
                  "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,
              ),
              (
                  "include parent dag",
                  "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,
              ),
              (
                  "include sub dag",
                  "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,
              ),
              (
                  "dry_run default",
                  "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,
              ),
          ]
      )
      @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_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,
              },
          ]
  
          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': 'sleep_for_0',
              },
              {
                  '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_1',
              },
              {
                  '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_2',
              },
              {
                  '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_3',
              },
          ]
          for task_instance in expected_response:
              assert task_instance in response.json["task_instances"]
          assert 5 == 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)
  
      @parameterized.expand(["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
  
      @parameterized.expand(
          [
              ({"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,
                  "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,
                  "include_upstream": True,
                  "include_downstream": True,
                  "include_future": True,
                  "include_past": True,
                  "new_state": "failed",
              },
          )
          assert_401(response)
  
      @parameterized.expand(["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,
                  "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,
                  "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,
                  "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,
                  "include_upstream": True,
                  "include_downstream": True,
                  "include_future": True,
                  "include_past": True,
                  "new_state": "failed",
              },
          )
          assert response.status_code == 404
  
      @parameterized.expand(
          [
              (
                  {
                      "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
  