text: # -*- coding: utf-8 -*-
  #
  # 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 multiprocessing
  import os
  import pickle  # type: ignore
  import signal
  import unittest
  from datetime import timedelta
  from time import sleep
  from unittest import mock
  
  import sqlalchemy
  from dateutil.relativedelta import relativedelta
  from numpy.testing import assert_array_almost_equal
  from pendulum import utcnow
  
  from airflow import DAG, configuration, exceptions, jobs, settings, utils
  from airflow.bin import cli
  from airflow.configuration import AirflowConfigException, conf, run_command
  from airflow.exceptions import AirflowException
  from airflow.executors import SequentialExecutor
  from airflow.hooks.base_hook import BaseHook
  from airflow.hooks.sqlite_hook import SqliteHook
  from airflow.models import Connection, DagBag, DagRun, Pool, TaskFail, TaskInstance, Variable
  from airflow.models.baseoperator import BaseOperator
  from airflow.operators.bash_operator import BashOperator
  from airflow.operators.check_operator import CheckOperator, ValueCheckOperator
  from airflow.operators.dummy_operator import DummyOperator
  from airflow.operators.python_operator import PythonOperator
  from airflow.settings import Session
  from airflow.utils import timezone
  from airflow.utils.dates import days_ago, infer_time_unit, round_time, scale_time_units
  from airflow.utils.state import State
  from airflow.utils.timezone import datetime
  from tests.test_utils.config import conf_vars
  
  DEV_NULL = '/dev/null'
  TEST_DAG_FOLDER = os.path.join(
      os.path.dirname(os.path.realpath(__file__)), 'dags')
  DEFAULT_DATE = datetime(2015, 1, 1)
  DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat()
  DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10]
  TEST_DAG_ID = 'unit_tests'
  EXAMPLE_DAG_DEFAULT_DATE = days_ago(2)
  
  
  class OperatorSubclass(BaseOperator):
      """
      An operator to test template substitution
      """
      template_fields = ['some_templated_field']
  
      def __init__(self, some_templated_field, *args, **kwargs):
          super().__init__(*args, **kwargs)
          self.some_templated_field = some_templated_field
  
      def execute(self, context):
          pass
  
  
  class TestCore(unittest.TestCase):
      TEST_SCHEDULE_WITH_NO_PREVIOUS_RUNS_DAG_ID = TEST_DAG_ID + 'test_schedule_dag_no_previous_runs'
      TEST_SCHEDULE_DAG_FAKE_SCHEDULED_PREVIOUS_DAG_ID = \
          TEST_DAG_ID + 'test_schedule_dag_fake_scheduled_previous'
      TEST_SCHEDULE_DAG_NO_END_DATE_UP_TO_TODAY_ONLY_DAG_ID = \
          TEST_DAG_ID + 'test_schedule_dag_no_end_date_up_to_today_only'
      TEST_SCHEDULE_ONCE_DAG_ID = TEST_DAG_ID + 'test_schedule_dag_once'
      TEST_SCHEDULE_RELATIVEDELTA_DAG_ID = TEST_DAG_ID + 'test_schedule_dag_relativedelta'
      TEST_SCHEDULE_START_END_DATES_DAG_ID = TEST_DAG_ID + 'test_schedule_dag_start_end_dates'
  
      default_scheduler_args = {"num_runs": 1}
  
      def setUp(self):
          self.dagbag = DagBag(
              dag_folder=DEV_NULL, include_examples=True)
          self.args = {'owner': 'airflow', 'start_date': DEFAULT_DATE}
          self.dag = DAG(TEST_DAG_ID, default_args=self.args)
          self.dag_bash = self.dagbag.dags['example_bash_operator']
          self.runme_0 = self.dag_bash.get_task('runme_0')
          self.run_after_loop = self.dag_bash.get_task('run_after_loop')
          self.run_this_last = self.dag_bash.get_task('run_this_last')
  
      def tearDown(self):
          if os.environ.get('KUBERNETES_VERSION') is not None:
              return
  
          dag_ids_to_clean = [
              TEST_DAG_ID,
              self.TEST_SCHEDULE_WITH_NO_PREVIOUS_RUNS_DAG_ID,
              self.TEST_SCHEDULE_DAG_FAKE_SCHEDULED_PREVIOUS_DAG_ID,
              self.TEST_SCHEDULE_DAG_NO_END_DATE_UP_TO_TODAY_ONLY_DAG_ID,
              self.TEST_SCHEDULE_ONCE_DAG_ID,
              self.TEST_SCHEDULE_RELATIVEDELTA_DAG_ID,
              self.TEST_SCHEDULE_START_END_DATES_DAG_ID,
          ]
          session = Session()
          session.query(DagRun).filter(
              DagRun.dag_id.in_(dag_ids_to_clean)).delete(
              synchronize_session=False)
          session.query(TaskInstance).filter(
              TaskInstance.dag_id.in_(dag_ids_to_clean)).delete(
              synchronize_session=False)
          session.query(TaskFail).filter(
              TaskFail.dag_id.in_(dag_ids_to_clean)).delete(
              synchronize_session=False)
          session.commit()
          session.close()
  
      def test_schedule_dag_no_previous_runs(self):
          """
          Tests scheduling a dag with no previous runs
          """
          dag = DAG(self.TEST_SCHEDULE_WITH_NO_PREVIOUS_RUNS_DAG_ID)
          dag.add_task(BaseOperator(
              task_id="faketastic",
              owner='Also fake',
              start_date=datetime(2015, 1, 2, 0, 0)))
  
          dag_run = jobs.SchedulerJob(**self.default_scheduler_args).create_dag_run(dag)
          self.assertIsNotNone(dag_run)
          self.assertEqual(dag.dag_id, dag_run.dag_id)
          self.assertIsNotNone(dag_run.run_id)
          self.assertNotEqual('', dag_run.run_id)
          self.assertEqual(
              datetime(2015, 1, 2, 0, 0),
              dag_run.execution_date,
              msg='dag_run.execution_date did not match expectation: {0}'
              .format(dag_run.execution_date)
          )
          self.assertEqual(State.RUNNING, dag_run.state)
          self.assertFalse(dag_run.external_trigger)
          dag.clear()
  
      def test_schedule_dag_relativedelta(self):
          """
          Tests scheduling a dag with a relativedelta schedule_interval
          """
          delta = relativedelta(hours=+1)
          dag = DAG(self.TEST_SCHEDULE_RELATIVEDELTA_DAG_ID,
                    schedule_interval=delta)
          dag.add_task(BaseOperator(
              task_id="faketastic",
              owner='Also fake',
              start_date=datetime(2015, 1, 2, 0, 0)))
  
          dag_run = jobs.SchedulerJob(**self.default_scheduler_args).create_dag_run(dag)
          self.assertIsNotNone(dag_run)
          self.assertEqual(dag.dag_id, dag_run.dag_id)
          self.assertIsNotNone(dag_run.run_id)
          self.assertNotEqual('', dag_run.run_id)
          self.assertEqual(
              datetime(2015, 1, 2, 0, 0),
              dag_run.execution_date,
              msg='dag_run.execution_date did not match expectation: {0}'
              .format(dag_run.execution_date)
          )
          self.assertEqual(State.RUNNING, dag_run.state)
          self.assertFalse(dag_run.external_trigger)
          dag_run2 = jobs.SchedulerJob(**self.default_scheduler_args).create_dag_run(dag)
          self.assertIsNotNone(dag_run2)
          self.assertEqual(dag.dag_id, dag_run2.dag_id)
          self.assertIsNotNone(dag_run2.run_id)
          self.assertNotEqual('', dag_run2.run_id)
          self.assertEqual(
              datetime(2015, 1, 2, 0, 0) + delta,
              dag_run2.execution_date,
              msg='dag_run2.execution_date did not match expectation: {0}'
              .format(dag_run2.execution_date)
          )
          self.assertEqual(State.RUNNING, dag_run2.state)
          self.assertFalse(dag_run2.external_trigger)
          dag.clear()
  
      def test_schedule_dag_fake_scheduled_previous(self):
          """
          Test scheduling a dag where there is a prior DagRun
          which has the same run_id as the next run should have
          """
          delta = timedelta(hours=1)
  
          dag = DAG(self.TEST_SCHEDULE_DAG_FAKE_SCHEDULED_PREVIOUS_DAG_ID,
                    schedule_interval=delta,
                    start_date=DEFAULT_DATE)
          dag.add_task(BaseOperator(
              task_id="faketastic",
              owner='Also fake',
              start_date=DEFAULT_DATE))
  
          scheduler = jobs.SchedulerJob(**self.default_scheduler_args)
          dag.create_dagrun(run_id=DagRun.id_for_date(DEFAULT_DATE),
                            execution_date=DEFAULT_DATE,
                            state=State.SUCCESS,
                            external_trigger=True)
          dag_run = scheduler.create_dag_run(dag)
          self.assertIsNotNone(dag_run)
          self.assertEqual(dag.dag_id, dag_run.dag_id)
          self.assertIsNotNone(dag_run.run_id)
          self.assertNotEqual('', dag_run.run_id)
          self.assertEqual(
              DEFAULT_DATE + delta,
              dag_run.execution_date,
              msg='dag_run.execution_date did not match expectation: {0}'
              .format(dag_run.execution_date)
          )
          self.assertEqual(State.RUNNING, dag_run.state)
          self.assertFalse(dag_run.external_trigger)
  
      def test_schedule_dag_once(self):
          """
          Tests scheduling a dag scheduled for @once - should be scheduled the first time
          it is called, and not scheduled the second.
          """
          dag = DAG(self.TEST_SCHEDULE_ONCE_DAG_ID)
          dag.schedule_interval = '@once'
          dag.add_task(BaseOperator(
              task_id="faketastic",
              owner='Also fake',
              start_date=datetime(2015, 1, 2, 0, 0)))
          dag_run = jobs.SchedulerJob(**self.default_scheduler_args).create_dag_run(dag)
          dag_run2 = jobs.SchedulerJob(**self.default_scheduler_args).create_dag_run(dag)
  
          self.assertIsNotNone(dag_run)
          self.assertIsNone(dag_run2)
          dag.clear()
  
      def test_fractional_seconds(self):
          """
          Tests if fractional seconds are stored in the database
          """
          dag = DAG(TEST_DAG_ID + 'test_fractional_seconds')
          dag.schedule_interval = '@once'
          dag.add_task(BaseOperator(
              task_id="faketastic",
              owner='Also fake',
              start_date=datetime(2015, 1, 2, 0, 0)))
  
          start_date = timezone.utcnow()
  
          run = dag.create_dagrun(
              run_id='test_' + start_date.isoformat(),
              execution_date=start_date,
              start_date=start_date,
              state=State.RUNNING,
              external_trigger=False
          )
  
          run.refresh_from_db()
  
          self.assertEqual(start_date, run.execution_date,
                           "dag run execution_date loses precision")
          self.assertEqual(start_date, run.start_date,
                           "dag run start_date loses precision ")
  
      def test_schedule_dag_start_end_dates(self):
          """
          Tests that an attempt to schedule a task after the Dag's end_date
          does not succeed.
          """
          delta = timedelta(hours=1)
          runs = 3
          start_date = DEFAULT_DATE
          end_date = start_date + (runs - 1) * delta
  
          dag = DAG(self.TEST_SCHEDULE_START_END_DATES_DAG_ID,
                    start_date=start_date,
                    end_date=end_date,
                    schedule_interval=delta)
          dag.add_task(BaseOperator(task_id='faketastic', owner='Also fake'))
  
          # Create and schedule the dag runs
          dag_runs = []
          scheduler = jobs.SchedulerJob(**self.default_scheduler_args)
          for _ in range(runs):
              dag_runs.append(scheduler.create_dag_run(dag))
  
          additional_dag_run = scheduler.create_dag_run(dag)
  
          for dag_run in dag_runs:
              self.assertIsNotNone(dag_run)
  
          self.assertIsNone(additional_dag_run)
  
      def test_schedule_dag_no_end_date_up_to_today_only(self):
          """
          Tests that a Dag created without an end_date can only be scheduled up
          to and including the current datetime.
  
          For example, if today is 2016-01-01 and we are scheduling from a
          start_date of 2015-01-01, only jobs up to, but not including
          2016-01-01 should be scheduled.
          """
          session = settings.Session()
          delta = timedelta(days=1)
          now = utcnow()
          start_date = now.subtract(weeks=1)
  
          runs = (now - start_date).days
  
          dag = DAG(self.TEST_SCHEDULE_DAG_NO_END_DATE_UP_TO_TODAY_ONLY_DAG_ID,
                    start_date=start_date,
                    schedule_interval=delta)
          dag.add_task(BaseOperator(task_id='faketastic', owner='Also fake'))
  
          dag_runs = []
          scheduler = jobs.SchedulerJob(**self.default_scheduler_args)
          for _ in range(runs):
              dag_run = scheduler.create_dag_run(dag)
              dag_runs.append(dag_run)
  
              # Mark the DagRun as complete
              dag_run.state = State.SUCCESS
              session.merge(dag_run)
              session.commit()
  
          # Attempt to schedule an additional dag run (for 2016-01-01)
          additional_dag_run = scheduler.create_dag_run(dag)
  
          for dag_run in dag_runs:
              self.assertIsNotNone(dag_run)
  
          self.assertIsNone(additional_dag_run)
  
      def test_confirm_unittest_mod(self):
          self.assertTrue(conf.get('core', 'unit_test_mode'))
  
      def test_pickling(self):
          dp = self.dag.pickle()
          self.assertEqual(dp.pickle.dag_id, self.dag.dag_id)
  
      def test_rich_comparison_ops(self):
  
          class DAGsubclass(DAG):
              pass
  
          dag_eq = DAG(TEST_DAG_ID, default_args=self.args)
  
          dag_diff_load_time = DAG(TEST_DAG_ID, default_args=self.args)
          dag_diff_name = DAG(TEST_DAG_ID + '_neq', default_args=self.args)
  
          dag_subclass = DAGsubclass(TEST_DAG_ID, default_args=self.args)
          dag_subclass_diff_name = DAGsubclass(
              TEST_DAG_ID + '2', default_args=self.args)
  
          for d in [dag_eq, dag_diff_name, dag_subclass, dag_subclass_diff_name]:
              d.last_loaded = self.dag.last_loaded
  
          # test identity equality
          self.assertEqual(self.dag, self.dag)
  
          # test dag (in)equality based on _comps
          self.assertEqual(dag_eq, self.dag)
          self.assertNotEqual(dag_diff_name, self.dag)
          self.assertNotEqual(dag_diff_load_time, self.dag)
  
          # test dag inequality based on type even if _comps happen to match
          self.assertNotEqual(dag_subclass, self.dag)
  
          # a dag should equal an unpickled version of itself
          d = pickle.dumps(self.dag)
          self.assertEqual(pickle.loads(d), self.dag)
  
          # dags are ordered based on dag_id no matter what the type is
          self.assertLess(self.dag, dag_diff_name)
          self.assertGreater(self.dag, dag_diff_load_time)
          self.assertLess(self.dag, dag_subclass_diff_name)
  
          # greater than should have been created automatically by functools
          self.assertGreater(dag_diff_name, self.dag)
  
          # hashes are non-random and match equality
          self.assertEqual(hash(self.dag), hash(self.dag))
          self.assertEqual(hash(dag_eq), hash(self.dag))
          self.assertNotEqual(hash(dag_diff_name), hash(self.dag))
          self.assertNotEqual(hash(dag_subclass), hash(self.dag))
  
      def test_check_operators(self):
  
          conn_id = "sqlite_default"
  
          captainHook = BaseHook.get_hook(conn_id=conn_id)
          captainHook.run("CREATE TABLE operator_test_table (a, b)")
          captainHook.run("insert into operator_test_table values (1,2)")
  
          t = CheckOperator(
              task_id='check',
              sql="select count(*) from operator_test_table",
              conn_id=conn_id,
              dag=self.dag)
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
          t = ValueCheckOperator(
              task_id='value_check',
              pass_value=95,
              tolerance=0.1,
              conn_id=conn_id,
              sql="SELECT 100",
              dag=self.dag)
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
          captainHook.run("drop table operator_test_table")
  
      def test_clear_api(self):
          task = self.dag_bash.tasks[0]
          task.clear(
              start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
              upstream=True, downstream=True)
          ti = TaskInstance(task=task, execution_date=DEFAULT_DATE)
          ti.are_dependents_done()
  
      def test_illegal_args(self):
          """
          Tests that Operators reject illegal arguments
          """
          msg = 'Invalid arguments were passed to BashOperator '
          '(task_id: test_illegal_args).'
          with conf_vars({('operators', 'allow_illegal_arguments'): 'True'}):
              with self.assertWarns(PendingDeprecationWarning) as warning:
                  BashOperator(
                      task_id='test_illegal_args',
                      bash_command='echo success',
                      dag=self.dag,
                      illegal_argument_1234='hello?')
                  assert any(msg in str(w) for w in warning.warnings)
  
      def test_illegal_args_forbidden(self):
          """
          Tests that operators raise exceptions on illegal arguments when
          illegal arguments are not allowed.
          """
          with self.assertRaises(AirflowException) as ctx:
              BashOperator(
                  task_id='test_illegal_args',
                  bash_command='echo success',
                  dag=self.dag,
                  illegal_argument_1234='hello?')
          self.assertIn(
              ('Invalid arguments were passed to BashOperator '
               '(task_id: test_illegal_args).'),
              str(ctx.exception))
  
      def test_bash_operator(self):
          t = BashOperator(
              task_id='test_bash_operator',
              bash_command="echo success",
              dag=self.dag)
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_bash_operator_multi_byte_output(self):
          t = BashOperator(
              task_id='test_multi_byte_bash_operator',
              bash_command="echo \u2600",
              dag=self.dag,
              output_encoding='utf-8')
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_bash_operator_kill(self):
          import psutil
          sleep_time = "100%d" % os.getpid()
          t = BashOperator(
              task_id='test_bash_operator_kill',
              execution_timeout=timedelta(seconds=1),
              bash_command="/bin/bash -c 'sleep %s'" % sleep_time,
              dag=self.dag)
          self.assertRaises(
              exceptions.AirflowTaskTimeout,
              t.run,
              start_date=DEFAULT_DATE, end_date=DEFAULT_DATE)
          sleep(2)
          pid = -1
          for proc in psutil.process_iter():
              if proc.cmdline() == ['sleep', sleep_time]:
                  pid = proc.pid
          if pid != -1:
              os.kill(pid, signal.SIGTERM)
              self.fail("BashOperator's subprocess still running after stopping on timeout!")
  
      def test_on_failure_callback(self):
          # Annoying workaround for nonlocal not existing in python 2
          data = {'called': False}
  
          def check_failure(context, test_case=self):
              data['called'] = True
              error = context.get('exception')
              test_case.assertIsInstance(error, AirflowException)
  
          t = BashOperator(
              task_id='check_on_failure_callback',
              bash_command="exit 1",
              dag=self.dag,
              on_failure_callback=check_failure)
          self.assertRaises(
              exceptions.AirflowException,
              t.run,
              start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
          self.assertTrue(data['called'])
  
      def test_dryrun(self):
          t = BashOperator(
              task_id='test_dryrun',
              bash_command="echo success",
              dag=self.dag)
          t.dry_run()
  
      def test_sqlite(self):
          import airflow.operators.sqlite_operator
          t = airflow.operators.sqlite_operator.SqliteOperator(
              task_id='time_sqlite',
              sql="CREATE TABLE IF NOT EXISTS unitest (dummy VARCHAR(20))",
              dag=self.dag)
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_timeout(self):
          t = PythonOperator(
              task_id='test_timeout',
              execution_timeout=timedelta(seconds=1),
              python_callable=lambda: sleep(5),
              dag=self.dag)
          self.assertRaises(
              exceptions.AirflowTaskTimeout,
              t.run,
              start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_python_op(self):
          def test_py_op(templates_dict, ds, **kwargs):
              if not templates_dict['ds'] == ds:
                  raise Exception("failure")
  
          t = PythonOperator(
              task_id='test_py_op',
              python_callable=test_py_op,
              templates_dict={'ds': "{{ ds }}"},
              dag=self.dag)
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_complex_template(self):
          def verify_templated_field(context):
              self.assertEqual(context['ti'].task.some_templated_field['bar'][1],
                               context['ds'])
  
          t = OperatorSubclass(
              task_id='test_complex_template',
              some_templated_field={
                  'foo': '123',
                  'bar': ['baz', '{{ ds }}']
              },
              dag=self.dag)
          t.execute = verify_templated_field
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_template_with_variable(self):
          """
          Test the availability of variables in templates
          """
          val = {
              'test_value': 'a test value'
          }
          Variable.set("a_variable", val['test_value'])
  
          def verify_templated_field(context):
              self.assertEqual(context['ti'].task.some_templated_field,
                               val['test_value'])
  
          t = OperatorSubclass(
              task_id='test_complex_template',
              some_templated_field='{{ var.value.a_variable }}',
              dag=self.dag)
          t.execute = verify_templated_field
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_template_with_json_variable(self):
          """
          Test the availability of variables (serialized as JSON) in templates
          """
          val = {
              'test_value': {'foo': 'bar', 'obj': {'v1': 'yes', 'v2': 'no'}}
          }
          Variable.set("a_variable", val['test_value'], serialize_json=True)
  
          def verify_templated_field(context):
              self.assertEqual(context['ti'].task.some_templated_field,
                               val['test_value']['obj']['v2'])
  
          t = OperatorSubclass(
              task_id='test_complex_template',
              some_templated_field='{{ var.json.a_variable.obj.v2 }}',
              dag=self.dag)
          t.execute = verify_templated_field
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_template_with_json_variable_as_value(self):
          """
          Test the availability of variables (serialized as JSON) in templates, but
          accessed as a value
          """
          val = {
              'test_value': {'foo': 'bar'}
          }
          Variable.set("a_variable", val['test_value'], serialize_json=True)
  
          def verify_templated_field(context):
              self.assertEqual(context['ti'].task.some_templated_field,
                               '{\n  "foo": "bar"\n}')
  
          t = OperatorSubclass(
              task_id='test_complex_template',
              some_templated_field='{{ var.value.a_variable }}',
              dag=self.dag)
          t.execute = verify_templated_field
          t.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
  
      def test_template_non_bool(self):
          """
          Test templates can handle objects with no sense of truthiness
          """
  
          class NonBoolObject:
              def __len__(self):
                  return NotImplemented
  
              def __bool__(self):
                  return NotImplemented
  
          t = OperatorSubclass(
              task_id='test_bad_template_obj',
              some_templated_field=NonBoolObject(),
              dag=self.dag)
          t.resolve_template_files()
  
      def test_task_get_template(self):
          TI = TaskInstance
          ti = TI(
              task=self.runme_0, execution_date=DEFAULT_DATE)
          ti.dag = self.dag_bash
          ti.run(ignore_ti_state=True)
          context = ti.get_template_context()
  
          # DEFAULT DATE is 2015-01-01
          self.assertEqual(context['ds'], '2015-01-01')
          self.assertEqual(context['ds_nodash'], '20150101')
  
          # next_ds is 2015-01-02 as the dag interval is daily
          self.assertEqual(context['next_ds'], '2015-01-02')
          self.assertEqual(context['next_ds_nodash'], '20150102')
  
          # prev_ds is 2014-12-31 as the dag interval is daily
          self.assertEqual(context['prev_ds'], '2014-12-31')
          self.assertEqual(context['prev_ds_nodash'], '20141231')
  
          self.assertEqual(context['ts'], '2015-01-01T00:00:00+00:00')
          self.assertEqual(context['ts_nodash'], '20150101T000000')
          self.assertEqual(context['ts_nodash_with_tz'], '20150101T000000+0000')
  
          self.assertEqual(context['yesterday_ds'], '2014-12-31')
          self.assertEqual(context['yesterday_ds_nodash'], '20141231')
  
          self.assertEqual(context['tomorrow_ds'], '2015-01-02')
          self.assertEqual(context['tomorrow_ds_nodash'], '20150102')
  
      def test_local_task_job(self):
          TI = TaskInstance
          ti = TI(
              task=self.runme_0, execution_date=DEFAULT_DATE)
          job = jobs.LocalTaskJob(task_instance=ti, ignore_ti_state=True)
          job.run()
  
      def test_raw_job(self):
          TI = TaskInstance
          ti = TI(
              task=self.runme_0, execution_date=DEFAULT_DATE)
          ti.dag = self.dag_bash
          ti.run(ignore_ti_state=True)
  
      def test_variable_set_get_round_trip(self):
          Variable.set("tested_var_set_id", "Monday morning breakfast")
          self.assertEqual("Monday morning breakfast", Variable.get("tested_var_set_id"))
  
      def test_variable_set_get_round_trip_json(self):
          value = {"a": 17, "b": 47}
          Variable.set("tested_var_set_id", value, serialize_json=True)
          self.assertEqual(value, Variable.get("tested_var_set_id", deserialize_json=True))
  
      def test_get_non_existing_var_should_return_default(self):
          default_value = "some default val"
          self.assertEqual(default_value, Variable.get("thisIdDoesNotExist",
                                                       default_var=default_value))
  
      def test_get_non_existing_var_should_raise_key_error(self):
          with self.assertRaises(KeyError):
              Variable.get("thisIdDoesNotExist")
  
      def test_get_non_existing_var_with_none_default_should_return_none(self):
          self.assertIsNone(Variable.get("thisIdDoesNotExist", default_var=None))
  
      def test_get_non_existing_var_should_not_deserialize_json_default(self):
          default_value = "}{ this is a non JSON default }{"
          self.assertEqual(default_value, Variable.get("thisIdDoesNotExist",
                                                       default_var=default_value,
                                                       deserialize_json=True))
  
      def test_variable_setdefault_round_trip(self):
          key = "tested_var_setdefault_1_id"
          value = "Monday morning breakfast in Paris"
          Variable.setdefault(key, value)
          self.assertEqual(value, Variable.get(key))
  
      def test_variable_setdefault_round_trip_json(self):
          key = "tested_var_setdefault_2_id"
          value = {"city": 'Paris', "Happiness": True}
          Variable.setdefault(key, value, deserialize_json=True)
          self.assertEqual(value, Variable.get(key, deserialize_json=True))
  
      def test_variable_setdefault_existing_json(self):
          key = "tested_var_setdefault_2_id"
          value = {"city": 'Paris', "Happiness": True}
          Variable.set(key, value, serialize_json=True)
          val = Variable.setdefault(key, value, deserialize_json=True)
          # Check the returned value, and the stored value are handled correctly.
          self.assertEqual(value, val)
          self.assertEqual(value, Variable.get(key, deserialize_json=True))
  
      def test_variable_delete(self):
          key = "tested_var_delete"
          value = "to be deleted"
  
          # No-op if the variable doesn't exist
          Variable.delete(key)
          with self.assertRaises(KeyError):
              Variable.get(key)
  
          # Set the variable
          Variable.set(key, value)
          self.assertEqual(value, Variable.get(key))
  
          # Delete the variable
          Variable.delete(key)
          with self.assertRaises(KeyError):
              Variable.get(key)
  
      def test_parameterized_config_gen(self):
  
          cfg = configuration.parameterized_config(configuration.DEFAULT_CONFIG)
  
          # making sure some basic building blocks are present:
          self.assertIn("[core]", cfg)
          self.assertIn("dags_folder", cfg)
          self.assertIn("sql_alchemy_conn", cfg)
          self.assertIn("fernet_key", cfg)
  
          # making sure replacement actually happened
          self.assertNotIn("{AIRFLOW_HOME}", cfg)
          self.assertNotIn("{FERNET_KEY}", cfg)
  
      def test_config_use_original_when_original_and_fallback_are_present(self):
          self.assertTrue(conf.has_option("core", "FERNET_KEY"))
          self.assertFalse(conf.has_option("core", "FERNET_KEY_CMD"))
  
          FERNET_KEY = conf.get('core', 'FERNET_KEY')
  
          with conf_vars({('core', 'FERNET_KEY_CMD'): 'printf HELLO'}):
              FALLBACK_FERNET_KEY = conf.get(
                  "core",
                  "FERNET_KEY"
              )
  
          self.assertEqual(FERNET_KEY, FALLBACK_FERNET_KEY)
  
      def test_config_throw_error_when_original_and_fallback_is_absent(self):
          self.assertTrue(conf.has_option("core", "FERNET_KEY"))
          self.assertFalse(conf.has_option("core", "FERNET_KEY_CMD"))
  
          with conf_vars({('core', 'fernet_key'): None}):
              with self.assertRaises(AirflowConfigException) as cm:
                  conf.get("core", "FERNET_KEY")
  
          exception = str(cm.exception)
          message = "section/key [core/fernet_key] not found in config"
          self.assertEqual(message, exception)
  
      def test_config_override_original_when_non_empty_envvar_is_provided(self):
          key = "AIRFLOW__CORE__FERNET_KEY"
          value = "some value"
  
          with mock.patch.dict('os.environ', {key: value}):
              FERNET_KEY = conf.get('core', 'FERNET_KEY')
  
          self.assertEqual(value, FERNET_KEY)
  
      def test_config_override_original_when_empty_envvar_is_provided(self):
          key = "AIRFLOW__CORE__FERNET_KEY"
          value = ""
  
          with mock.patch.dict('os.environ', {key: value}):
              FERNET_KEY = conf.get('core', 'FERNET_KEY')
  
          self.assertEqual(value, FERNET_KEY)
  
      def test_round_time(self):
  
          rt1 = round_time(datetime(2015, 1, 1, 6), timedelta(days=1))
          self.assertEqual(datetime(2015, 1, 1, 0, 0), rt1)
  
          rt2 = round_time(datetime(2015, 1, 2), relativedelta(months=1))
          self.assertEqual(datetime(2015, 1, 1, 0, 0), rt2)
  
          rt3 = round_time(datetime(2015, 9, 16, 0, 0), timedelta(1), datetime(
              2015, 9, 14, 0, 0))
          self.assertEqual(datetime(2015, 9, 16, 0, 0), rt3)
  
          rt4 = round_time(datetime(2015, 9, 15, 0, 0), timedelta(1), datetime(
              2015, 9, 14, 0, 0))
          self.assertEqual(datetime(2015, 9, 15, 0, 0), rt4)
  
          rt5 = round_time(datetime(2015, 9, 14, 0, 0), timedelta(1), datetime(
              2015, 9, 14, 0, 0))
          self.assertEqual(datetime(2015, 9, 14, 0, 0), rt5)
  
          rt6 = round_time(datetime(2015, 9, 13, 0, 0), timedelta(1), datetime(
              2015, 9, 14, 0, 0))
          self.assertEqual(datetime(2015, 9, 14, 0, 0), rt6)
  
      def test_infer_time_unit(self):
  
          self.assertEqual('minutes', infer_time_unit([130, 5400, 10]))
  
          self.assertEqual('seconds', infer_time_unit([110, 50, 10, 100]))
  
          self.assertEqual('hours', infer_time_unit([100000, 50000, 10000, 20000]))
  
          self.assertEqual('days', infer_time_unit([200000, 100000]))
  
      def test_scale_time_units(self):
  
          # use assert_almost_equal from numpy.testing since we are comparing
          # floating point arrays
          arr1 = scale_time_units([130, 5400, 10], 'minutes')
          assert_array_almost_equal(arr1, [2.167, 90.0, 0.167], decimal=3)
  
          arr2 = scale_time_units([110, 50, 10, 100], 'seconds')
          assert_array_almost_equal(arr2, [110.0, 50.0, 10.0, 100.0], decimal=3)
  
          arr3 = scale_time_units([100000, 50000, 10000, 20000], 'hours')
          assert_array_almost_equal(arr3, [27.778, 13.889, 2.778, 5.556],
                                    decimal=3)
  
          arr4 = scale_time_units([200000, 100000], 'days')
          assert_array_almost_equal(arr4, [2.315, 1.157], decimal=3)
  
      def test_bad_trigger_rule(self):
          with self.assertRaises(AirflowException):
              DummyOperator(
                  task_id='test_bad_trigger',
                  trigger_rule="non_existent",
                  dag=self.dag)
  
      def test_terminate_task(self):
          """If a task instance's db state get deleted, it should fail"""
          TI = TaskInstance
          dag = self.dagbag.dags.get('test_utils')
          task = dag.task_dict.get('sleeps_forever')
  
          ti = TI(task=task, execution_date=DEFAULT_DATE)
          job = jobs.LocalTaskJob(
              task_instance=ti, ignore_ti_state=True, executor=SequentialExecutor())
  
          # Running task instance asynchronously
          p = multiprocessing.Process(target=job.run)
          p.start()
          sleep(5)
          settings.engine.dispose()
          session = settings.Session()
          ti.refresh_from_db(session=session)
          # making sure it's actually running
          self.assertEqual(State.RUNNING, ti.state)
          ti = session.query(TI).filter_by(
              dag_id=task.dag_id,
              task_id=task.task_id,
              execution_date=DEFAULT_DATE
          ).one()
  
          # deleting the instance should result in a failure
          session.delete(ti)
          session.commit()
          # waiting for the async task to finish
          p.join()
  
          # making sure that the task ended up as failed
          ti.refresh_from_db(session=session)
          self.assertEqual(State.FAILED, ti.state)
          session.close()
  
      def test_task_fail_duration(self):
          """If a task fails, the duration should be recorded in TaskFail"""
  
          p = BashOperator(
              task_id='pass_sleepy',
              bash_command='sleep 3',
              dag=self.dag)
          f = BashOperator(
              task_id='fail_sleepy',
              bash_command='sleep 5',
              execution_timeout=timedelta(seconds=3),
              retry_delay=timedelta(seconds=0),
              dag=self.dag)
          session = settings.Session()
          try:
              p.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
          except Exception:
              pass
          try:
              f.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
          except Exception:
              pass
          p_fails = session.query(TaskFail).filter_by(
              task_id='pass_sleepy',
              dag_id=self.dag.dag_id,
              execution_date=DEFAULT_DATE).all()
          f_fails = session.query(TaskFail).filter_by(
              task_id='fail_sleepy',
              dag_id=self.dag.dag_id,
              execution_date=DEFAULT_DATE).all()
  
          self.assertEqual(0, len(p_fails))
          self.assertEqual(1, len(f_fails))
          self.assertGreaterEqual(sum([f.duration for f in f_fails]), 3)
  
      def test_run_command(self):
          write = r'sys.stdout.buffer.write("\u1000foo".encode("utf8"))'
  
          cmd = 'import sys; {0}; sys.stdout.flush()'.format(write)
  
          self.assertEqual(run_command("python -c '{0}'".format(cmd)), '\u1000foo')
  
          self.assertEqual(run_command('echo "foo bar"'), 'foo bar\n')
          self.assertRaises(AirflowConfigException, run_command, 'bash -c "exit 1"')
  
      def test_externally_triggered_dagrun(self):
          TI = TaskInstance
  
          # Create the dagrun between two "scheduled" execution dates of the DAG
          EXECUTION_DATE = DEFAULT_DATE + timedelta(days=2)
          EXECUTION_DS = EXECUTION_DATE.strftime('%Y-%m-%d')
          EXECUTION_DS_NODASH = EXECUTION_DS.replace('-', '')
  
          dag = DAG(
              TEST_DAG_ID,
              default_args=self.args,
              schedule_interval=timedelta(weeks=1),
              start_date=DEFAULT_DATE)
          task = DummyOperator(task_id='test_externally_triggered_dag_context',
                               dag=dag)
          dag.create_dagrun(run_id=DagRun.id_for_date(EXECUTION_DATE),
                            execution_date=EXECUTION_DATE,
                            state=State.RUNNING,
                            external_trigger=True)
          task.run(
              start_date=EXECUTION_DATE, end_date=EXECUTION_DATE)
  
          ti = TI(task=task, execution_date=EXECUTION_DATE)
          context = ti.get_template_context()
  
          # next_ds/prev_ds should be the execution date for manually triggered runs
          self.assertEqual(context['next_ds'], EXECUTION_DS)
          self.assertEqual(context['next_ds_nodash'], EXECUTION_DS_NODASH)
  
          self.assertEqual(context['prev_ds'], EXECUTION_DS)
          self.assertEqual(context['prev_ds_nodash'], EXECUTION_DS_NODASH)
  
  
  class TestCli(unittest.TestCase):
  
      TEST_USER1_EMAIL = 'test-user1@example.com'
      TEST_USER2_EMAIL = 'test-user2@example.com'
  
      @classmethod
      def setUpClass(cls):
          super().setUpClass()
          cls._cleanup()
  
      def setUp(self):
          super().setUp()
          from airflow.www import app as application
          self.app, self.appbuilder = application.create_app(session=Session, testing=True)
          self.app.config['TESTING'] = True
  
          self.parser = cli.CLIFactory.get_parser()
          self.dagbag = DagBag(dag_folder=DEV_NULL, include_examples=True)
          settings.configure_orm()
          self.session = Session
  
      def tearDown(self):
          self._cleanup(session=self.session)
          for email in [self.TEST_USER1_EMAIL, self.TEST_USER2_EMAIL]:
              test_user = self.appbuilder.sm.find_user(email=email)
              if test_user:
                  self.appbuilder.sm.del_register_user(test_user)
          for role_name in ['FakeTeamA', 'FakeTeamB']:
              if self.appbuilder.sm.find_role(role_name):
                  self.appbuilder.sm.delete_role(role_name)
  
          super().tearDown()
  
      @staticmethod
      def _cleanup(session=None):
          if session is None:
              session = Session()
  
          session.query(Pool).delete()
          session.query(Variable).delete()
          session.commit()
          session.close()
  
  
  class TestConnection(unittest.TestCase):
      def setUp(self):
          utils.db.initdb()
  
      @mock.patch.dict('os.environ', {
          'AIRFLOW_CONN_TEST_URI': 'postgres://username:password@ec2.compute.com:5432/the_database',
      })
      def test_using_env_var(self):
          c = SqliteHook.get_connection(conn_id='test_uri')
          self.assertEqual('ec2.compute.com', c.host)
          self.assertEqual('the_database', c.schema)
          self.assertEqual('username', c.login)
          self.assertEqual('password', c.password)
          self.assertEqual(5432, c.port)
  
      @mock.patch.dict('os.environ', {
          'AIRFLOW_CONN_TEST_URI_NO_CREDS': 'postgres://ec2.compute.com/the_database',
      })
      def test_using_unix_socket_env_var(self):
          c = SqliteHook.get_connection(conn_id='test_uri_no_creds')
          self.assertEqual('ec2.compute.com', c.host)
          self.assertEqual('the_database', c.schema)
          self.assertIsNone(c.login)
          self.assertIsNone(c.password)
          self.assertIsNone(c.port)
  
      def test_param_setup(self):
          c = Connection(conn_id='local_mysql', conn_type='mysql',
                         host='localhost', login='airflow',
                         password='airflow', schema='airflow')
          self.assertEqual('localhost', c.host)
          self.assertEqual('airflow', c.schema)
          self.assertEqual('airflow', c.login)
          self.assertEqual('airflow', c.password)
          self.assertIsNone(c.port)
  
      def test_env_var_priority(self):
          c = SqliteHook.get_connection(conn_id='airflow_db')
          self.assertNotEqual('ec2.compute.com', c.host)
  
          with mock.patch.dict('os.environ', {
              'AIRFLOW_CONN_AIRFLOW_DB': 'postgres://username:password@ec2.compute.com:5432/the_database',
          }):
              c = SqliteHook.get_connection(conn_id='airflow_db')
              self.assertEqual('ec2.compute.com', c.host)
              self.assertEqual('the_database', c.schema)
              self.assertEqual('username', c.login)
              self.assertEqual('password', c.password)
              self.assertEqual(5432, c.port)
  
      @mock.patch.dict('os.environ', {
          'AIRFLOW_CONN_TEST_URI': 'postgres://username:password@ec2.compute.com:5432/the_database',
          'AIRFLOW_CONN_TEST_URI_NO_CREDS': 'postgres://ec2.compute.com/the_database',
      })
      def test_dbapi_get_uri(self):
          conn = BaseHook.get_connection(conn_id='test_uri')
          hook = conn.get_hook()
          self.assertEqual('postgres://username:password@ec2.compute.com:5432/the_database', hook.get_uri())
          conn2 = BaseHook.get_connection(conn_id='test_uri_no_creds')
          hook2 = conn2.get_hook()
          self.assertEqual('postgres://ec2.compute.com/the_database', hook2.get_uri())
  
      @mock.patch.dict('os.environ', {
          'AIRFLOW_CONN_TEST_URI': 'postgres://username:password@ec2.compute.com:5432/the_database',
          'AIRFLOW_CONN_TEST_URI_NO_CREDS': 'postgres://ec2.compute.com/the_database',
      })
      def test_dbapi_get_sqlalchemy_engine(self):
          conn = BaseHook.get_connection(conn_id='test_uri')
          hook = conn.get_hook()
          engine = hook.get_sqlalchemy_engine()
          self.assertIsInstance(engine, sqlalchemy.engine.Engine)
          self.assertEqual('postgres://username:password@ec2.compute.com:5432/the_database', str(engine.url))
  
      @mock.patch.dict('os.environ', {
          'AIRFLOW_CONN_TEST_URI': 'postgres://username:password@ec2.compute.com:5432/the_database',
          'AIRFLOW_CONN_TEST_URI_NO_CREDS': 'postgres://ec2.compute.com/the_database',
      })
      def test_get_connections_env_var(self):
          conns = SqliteHook.get_connections(conn_id='test_uri')
          assert len(conns) == 1
          assert conns[0].host == 'ec2.compute.com'
          assert conns[0].schema == 'the_database'
          assert conns[0].login == 'username'
          assert conns[0].password == 'password'
          assert conns[0].port == 5432
  
  
  if __name__ == '__main__':
      unittest.main()
  