Python pendulum.datetime() Examples

def test_refresh_from_task(pool_override):
    task = DummyOperator(task_id="dummy", queue="test_queue", pool="test_pool1", pool_slots=3,
                         priority_weight=10, run_as_user="test", retries=30,
                         executor_config={"KubernetesExecutor": {"image": "myCustomDockerImage"}})
    ti = TI(task, execution_date=pendulum.datetime(2020, 1, 1))
    ti.refresh_from_task(task, pool_override=pool_override)

    assert ti.queue == task.queue

    if pool_override:
        assert ti.pool == pool_override
    else:
        assert ti.pool == task.pool

    assert ti.pool_slots == task.pool_slots
    assert ti.priority_weight == task.priority_weight_total
    assert ti.run_as_user == task.run_as_user
    assert ti.max_tries == task.retries
    assert ti.executor_config == task.executor_config
    assert ti.operator == DummyOperator.__name__ 

def test_ti_updates_with_task(self, session=None):
        """
        test that updating the executor_config propogates to the TaskInstance DB
        """
        with models.DAG(dag_id='test_run_pooling_task') as dag:
            task = DummyOperator(task_id='test_run_pooling_task_op', owner='airflow',
                                 executor_config={'foo': 'bar'},
                                 start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
        ti = TI(
            task=task, execution_date=timezone.utcnow())

        ti.run(session=session)
        tis = dag.get_task_instances()
        self.assertEqual({'foo': 'bar'}, tis[0].executor_config)
        with models.DAG(dag_id='test_run_pooling_task') as dag:
            task2 = DummyOperator(task_id='test_run_pooling_task_op', owner='airflow',
                                  executor_config={'bar': 'baz'},
                                  start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))

        ti = TI(
            task=task2, execution_date=timezone.utcnow())
        ti.run(session=session)
        tis = dag.get_task_instances()
        self.assertEqual({'bar': 'baz'}, tis[1].executor_config) 

def test_run_pooling_task_with_mark_success(self):
        """
        test that running task in an existing pool with mark_success param
        update task state as SUCCESS without running task
        despite it fails dependency checks.
        """
        dag = models.DAG(dag_id='test_run_pooling_task_with_mark_success')
        task = DummyOperator(
            task_id='test_run_pooling_task_with_mark_success_op',
            dag=dag,
            pool='test_pool',
            owner='airflow',
            start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
        ti = TI(
            task=task, execution_date=timezone.utcnow())
        ti.run(mark_success=True)
        self.assertEqual(ti.state, State.SUCCESS) 

def test_from_human_with_timezone(self):
        working_hours = WeeklyInterval('Mon', time(9, 00), 'Fri', time(17, 0))

        # During summer time
        now = pendulum.datetime(2019, 8, 30, tz='Europe/London')
        set_test_now(now)
        assert WeeklyInterval.from_human(
               "Mon 10:00 Europe/London - Fri 18:00 Europe/London"
        ) == working_hours

        # Outside summer time
        now = pendulum.datetime(2019, 12, 30, tz='Europe/London')
        set_test_now(now)
        assert WeeklyInterval.from_human(
            "Mon 09:00 Europe/London - Fri 17:00 Europe/London"
        ) == working_hours 

def test_mark_non_runnable_task_as_success(self):
        """
        test that running task with mark_success param update task state
        as SUCCESS without running task despite it fails dependency checks.
        """
        non_runnable_state = (
            set(State.task_states) - RUNNABLE_STATES - set(State.SUCCESS)).pop()
        dag = models.DAG(dag_id='test_mark_non_runnable_task_as_success')
        task = DummyOperator(
            task_id='test_mark_non_runnable_task_as_success_op',
            dag=dag,
            pool='test_pool',
            owner='airflow',
            start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
        ti = TI(
            task=task, execution_date=timezone.utcnow(), state=non_runnable_state)
        # TI.run() will sync from DB before validating deps.
        with create_session() as session:
            session.add(ti)
            session.commit()
        ti.run(mark_success=True)
        self.assertEqual(ti.state, State.SUCCESS) 

def test_next_retry_datetime_short_intervals(self):
        delay = datetime.timedelta(seconds=1)
        max_delay = datetime.timedelta(minutes=60)

        dag = models.DAG(dag_id='fail_dag')
        task = BashOperator(
            task_id='task_with_exp_backoff_and_short_time_interval',
            bash_command='exit 1',
            retries=3,
            retry_delay=delay,
            retry_exponential_backoff=True,
            max_retry_delay=max_delay,
            dag=dag,
            owner='airflow',
            start_date=timezone.datetime(2016, 2, 1, 0, 0, 0))
        ti = TI(
            task=task, execution_date=DEFAULT_DATE)
        ti.end_date = pendulum.instance(timezone.utcnow())

        date = ti.next_retry_datetime()
        # between 1 * 2^0.5 and 1 * 2^1 (15 and 30)
        period = ti.end_date.add(seconds=15) - ti.end_date.add(seconds=1)
        self.assertTrue(date in period) 

def test_from_human_with_timezone(self):
        weekly_1 = WeeklyInterval('Mon', time(10, 00), 'Fri', time(18, 00))
        weekly_2 = WeeklyInterval('Sat', time(12, 00), 'Sun', time(9, 00))
        interval = IntervalUnion([weekly_1, weekly_2])

        # During summer time
        now = pendulum.datetime(2019, 8, 30, tz='Europe/London')
        set_test_now(now)
        assert IntervalUnion.from_human(
            "Mon 11:00 Europe/London - Fri 19:00 Europe/London,"
            "Sat 13:00 Europe/London - Sun 10:00 Europe/London"
        ) == interval

        # Outside summer time
        now = pendulum.datetime(2019, 12, 30, tz='Europe/London')
        set_test_now(now)
        assert IntervalUnion.from_human(
            "Mon 10:00 Europe/London - Fri 18:00 Europe/London,"
            "Sat 12:00 Europe/London - Sun 09:00 Europe/London"
        ) == interval 

def test_xcom_push_flag(self):
        """
        Tests the option for Operators to push XComs
        """
        value = 'hello'
        task_id = 'test_no_xcom_push'
        dag = models.DAG(dag_id='test_xcom')

        # nothing saved to XCom
        task = PythonOperator(
            task_id=task_id,
            dag=dag,
            python_callable=lambda: value,
            do_xcom_push=False,
            owner='airflow',
            start_date=datetime.datetime(2017, 1, 1)
        )
        ti = TI(task=task, execution_date=datetime.datetime(2017, 1, 1))
        ti.run()
        self.assertEqual(
            ti.xcom_pull(
                task_ids=task_id, key=models.XCOM_RETURN_KEY
            ),
            None
        ) 

def test_post_execute_hook(self):
        """
        Test that post_execute hook is called with the Operator's result.
        The result ('error') will cause an error to be raised and trapped.
        """

        class TestError(Exception):
            pass

        class TestOperator(PythonOperator):
            def post_execute(self, context, result=None):
                if result == 'error':
                    raise TestError('expected error.')

        dag = models.DAG(dag_id='test_post_execute_dag')
        task = TestOperator(
            task_id='test_operator',
            dag=dag,
            python_callable=lambda: 'error',
            owner='airflow',
            start_date=timezone.datetime(2017, 2, 1))
        ti = TI(task=task, execution_date=timezone.utcnow())

        with self.assertRaises(TestError):
            ti.run() 

def test_email_alert(self, mock_send_email):
        dag = models.DAG(dag_id='test_failure_email')
        task = BashOperator(
            task_id='test_email_alert',
            dag=dag,
            bash_command='exit 1',
            start_date=DEFAULT_DATE,
            email='to')

        ti = TI(task=task, execution_date=datetime.datetime.now())

        try:
            ti.run()
        except AirflowException:
            pass

        (email, title, body), _ = mock_send_email.call_args
        self.assertEqual(email, 'to')
        self.assertIn('test_email_alert', title)
        self.assertIn('test_email_alert', body)
        self.assertIn('Try 1', body) 

def test_email_alert_with_config(self, mock_send_email):
        dag = models.DAG(dag_id='test_failure_email')
        task = BashOperator(
            task_id='test_email_alert_with_config',
            dag=dag,
            bash_command='exit 1',
            start_date=DEFAULT_DATE,
            email='to')

        ti = TI(
            task=task, execution_date=datetime.datetime.now())

        opener = mock_open(read_data='template: {{ti.task_id}}')
        with patch('airflow.models.taskinstance.open', opener, create=True):
            try:
                ti.run()
            except AirflowException:
                pass

        (email, title, body), _ = mock_send_email.call_args
        self.assertEqual(email, 'to')
        self.assertEqual('template: test_email_alert_with_config', title)
        self.assertEqual('template: test_email_alert_with_config', body) 

def test_execute_callback(self):
        called = False

        def on_execute_callable(context):
            nonlocal called
            called = True
            self.assertEqual(
                context['dag_run'].dag_id,
                'test_dagrun_execute_callback'
            )

        dag = DAG('test_execute_callback', start_date=DEFAULT_DATE,
                  end_date=DEFAULT_DATE + datetime.timedelta(days=10))
        task = DummyOperator(task_id='op', email='test@test.test',
                             on_execute_callback=on_execute_callable,
                             dag=dag)
        ti = TI(task=task, execution_date=datetime.datetime.now())
        ti.state = State.RUNNING
        session = settings.Session()
        session.merge(ti)
        session.commit()
        ti._run_raw_task()
        assert called
        ti.refresh_from_db()
        assert ti.state == State.SUCCESS 

def test_does_not_retry_on_airflow_fail_exception(self):
        def fail():
            raise AirflowFailException("hopeless")

        dag = models.DAG(dag_id='test_does_not_retry_on_airflow_fail_exception')
        task = PythonOperator(
            task_id='test_raise_airflow_fail_exception',
            dag=dag,
            python_callable=fail,
            owner='airflow',
            start_date=timezone.datetime(2016, 2, 1, 0, 0, 0),
            retries=1
        )
        ti = TI(task=task, execution_date=timezone.utcnow())
        try:
            ti.run()
        except AirflowFailException:
            pass  # expected
        self.assertEqual(State.FAILED, ti.state) 

def test_arbitrary_schedule():
    sample_date = pendulum.datetime(2019, 1, 31, 12, tz="America/Los_Angeles")
    schedule_path = "../sample_data/schedule/804_lametro-rail/2019-01-31.csv"

    assert (
        get_appropriate_timetable(
            sample_date, f"../sample_data/schedule/{line}_{agency}"
        )["path"]
        == schedule_path
    )

    assert (
        get_date_if_exists_otherwise_previous(
            sample_date, f"../sample_data/schedule/{line}_{agency}"
        )
        == sample_date
    ) 

def test_task_stats(self, stats_mock):
        dag = DAG('test_task_start_end_stats', start_date=DEFAULT_DATE,
                  end_date=DEFAULT_DATE + datetime.timedelta(days=10))
        op = DummyOperator(task_id='dummy_op', dag=dag)
        dag.create_dagrun(
            run_id='manual__' + DEFAULT_DATE.isoformat(),
            execution_date=DEFAULT_DATE,
            start_date=DEFAULT_DATE,
            state=State.RUNNING,
            external_trigger=False)
        ti = TI(task=op, execution_date=DEFAULT_DATE)
        ti.state = State.RUNNING
        session = settings.Session()
        session.merge(ti)
        session.commit()
        ti._run_raw_task()
        ti.refresh_from_db()
        stats_mock.assert_called_with('ti.finish.{}.{}.{}'.format(dag.dag_id, op.task_id, ti.state))
        self.assertIn(call('ti.start.{}.{}'.format(dag.dag_id, op.task_id)), stats_mock.mock_calls)
        self.assertEqual(stats_mock.call_count, 5) 

def test_echo_env_variables(self):
        dag = DAG('test_echo_env_variables', start_date=DEFAULT_DATE,
                  end_date=DEFAULT_DATE + datetime.timedelta(days=10))
        op = PythonOperator(task_id='hive_in_python_op',
                            dag=dag,
                            python_callable=self._env_var_check_callback)
        dag.create_dagrun(
            run_type=DagRunType.MANUAL,
            execution_date=DEFAULT_DATE,
            start_date=DEFAULT_DATE,
            state=State.RUNNING,
            external_trigger=False)
        ti = TI(task=op, execution_date=DEFAULT_DATE)
        ti.state = State.RUNNING
        session = settings.Session()
        session.merge(ti)
        session.commit()
        ti._run_raw_task()
        ti.refresh_from_db()
        self.assertEqual(ti.state, State.SUCCESS) 

def test_no_skipmixin_parent():
    """
    A simple DAG with no branching. Both op1 and op2 are DummyOperator. NotPreviouslySkippedDep is met.
    """
    start_date = pendulum.datetime(2020, 1, 1)
    dag = DAG(
        "test_no_skipmixin_parent_dag", schedule_interval=None, start_date=start_date
    )
    op1 = DummyOperator(task_id="op1", dag=dag)
    op2 = DummyOperator(task_id="op2", dag=dag)
    op1 >> op2

    ti2 = TaskInstance(op2, start_date)

    with create_session() as session:
        dep = NotPreviouslySkippedDep()
        assert len(list(dep.get_dep_statuses(ti2, session, DepContext()))) == 0
        assert dep.is_met(ti2, session)
        assert ti2.state != State.SKIPPED 

def test_parent_follow_branch():
    """
    A simple DAG with a BranchPythonOperator that follows op2. NotPreviouslySkippedDep is met.
    """
    start_date = pendulum.datetime(2020, 1, 1)
    dag = DAG(
        "test_parent_follow_branch_dag", schedule_interval=None, start_date=start_date
    )
    op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op2", dag=dag)
    op2 = DummyOperator(task_id="op2", dag=dag)
    op1 >> op2

    TaskInstance(op1, start_date).run()
    ti2 = TaskInstance(op2, start_date)

    with create_session() as session:
        dep = NotPreviouslySkippedDep()
        assert len(list(dep.get_dep_statuses(ti2, session, DepContext()))) == 0
        assert dep.is_met(ti2, session)
        assert ti2.state != State.SKIPPED 

def test_parent_skip_branch():
    """
    A simple DAG with a BranchPythonOperator that does not follow op2. NotPreviouslySkippedDep is not met.
    """
    start_date = pendulum.datetime(2020, 1, 1)
    dag = DAG(
        "test_parent_skip_branch_dag", schedule_interval=None, start_date=start_date
    )
    op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3", dag=dag)
    op2 = DummyOperator(task_id="op2", dag=dag)
    op3 = DummyOperator(task_id="op3", dag=dag)
    op1 >> [op2, op3]

    TaskInstance(op1, start_date).run()
    ti2 = TaskInstance(op2, start_date)

    with create_session() as session:
        dep = NotPreviouslySkippedDep()
        assert len(list(dep.get_dep_statuses(ti2, session, DepContext()))) == 1
        assert not dep.is_met(ti2, session)
        assert ti2.state == State.SKIPPED 

def test_parent_not_executed():
    """
    A simple DAG with a BranchPythonOperator that does not follow op2. Parent task is not yet
    executed (no xcom data). NotPreviouslySkippedDep is met (no decision).
    """
    start_date = pendulum.datetime(2020, 1, 1)
    dag = DAG(
        "test_parent_not_executed_dag", schedule_interval=None, start_date=start_date
    )
    op1 = BranchPythonOperator(task_id="op1", python_callable=lambda: "op3", dag=dag)
    op2 = DummyOperator(task_id="op2", dag=dag)
    op3 = DummyOperator(task_id="op3", dag=dag)
    op1 >> [op2, op3]

    ti2 = TaskInstance(op2, start_date)

    with create_session() as session:
        dep = NotPreviouslySkippedDep()
        assert len(list(dep.get_dep_statuses(ti2, session, DepContext()))) == 0
        assert dep.is_met(ti2, session)
        assert ti2.state == State.NONE 

def remount_workspace_ts(self, workspace_fs):
        date, time = TimestampedWorkspaceWidget.exec_modal(
            workspace_fs=workspace_fs, jobs_ctx=self.jobs_ctx, parent=self
        )

        if not date or not time:
            return

        datetime = pendulum.datetime(
            date.year(),
            date.month(),
            date.day(),
            time.hour(),
            time.minute(),
            time.second(),
            tzinfo="local",
        )
        self.mount_workspace(workspace_fs.workspace_id, datetime) 

def __new__(
            cls,
            year,
            month,
            day,
            hour=0,
            minute=0,
            second=0,
            microsecond=0,
            tzinfo=None,
            fold=0,
        ):
            self = datetime.datetime.__new__(
                cls, year, month, day, hour, minute, second, microsecond, tzinfo=tzinfo
            )

            self._fold = fold

            return self 

def timestamp(self):
            if self.tzinfo is None:
                s = timestamp(self)

                return s + self.microsecond / 1e6
            else:
                kwargs = {"tzinfo": self.tzinfo}

                if _HAS_FOLD:
                    kwargs["fold"] = self.fold

                dt = datetime.datetime(
                    self.year,
                    self.month,
                    self.day,
                    self.hour,
                    self.minute,
                    self.second,
                    self.microsecond,
                    **kwargs
                )
                return (dt - self._EPOCH).total_seconds() 

def int_timestamp(self):
        # Workaround needed to avoid inaccuracy
        # for far into the future datetimes
        kwargs = {"tzinfo": self.tzinfo}

        if _HAS_FOLD:
            kwargs["fold"] = self.fold

        dt = datetime.datetime(
            self.year,
            self.month,
            self.day,
            self.hour,
            self.minute,
            self.second,
            self.microsecond,
            **kwargs
        )

        delta = dt - self._EPOCH

        return delta.days * SECONDS_PER_DAY + delta.seconds 

def closest(self, dt1, dt2, *dts):
        """
        Get the farthest date from the instance.

        :type dt1: datetime.datetime
        :type dt2: datetime.datetime
        :type dts: list[datetime.datetime,]

        :rtype: DateTime
        """
        dt1 = pendulum.instance(dt1)
        dt2 = pendulum.instance(dt2)
        dts = [dt1, dt2] + [pendulum.instance(x) for x in dts]
        dts = [(abs(self - dt), dt) for dt in dts]

        return min(dts)[1] 

def farthest(self, dt1, dt2, *dts):
        """
        Get the farthest date from the instance.

        :type dt1: datetime.datetime
        :type dt2: datetime.datetime
        :type dts: list[datetime.datetime,]

        :rtype: DateTime
        """
        dt1 = pendulum.instance(dt1)
        dt2 = pendulum.instance(dt2)

        dts = [dt1, dt2] + [pendulum.instance(x) for x in dts]
        dts = [(abs(self - dt), dt) for dt in dts]

        return max(dts)[1] 

def average(self, dt=None):
        """
        Modify the current instance to the average
        of a given instance (default now) and the current instance.

        :type dt: DateTime or datetime

        :rtype: DateTime
        """
        if dt is None:
            dt = self.now(self.tz)

        diff = self.diff(dt, False)
        return self.add(
            microseconds=(diff.in_seconds() * 1000000 + diff.microseconds) // 2
        ) 

def __sub__(self, other):
        if isinstance(other, datetime.timedelta):
            return self._subtract_timedelta(other)

        if not isinstance(other, datetime.datetime):
            return NotImplemented

        if not isinstance(other, self.__class__):
            if other.tzinfo is None:
                other = pendulum.naive(
                    other.year,
                    other.month,
                    other.day,
                    other.hour,
                    other.minute,
                    other.second,
                    other.microsecond,
                )
            else:
                other = pendulum.instance(other)

        return other.diff(self, False) 

def __rsub__(self, other):
        if not isinstance(other, datetime.datetime):
            return NotImplemented

        if not isinstance(other, self.__class__):
            if other.tzinfo is None:
                other = pendulum.naive(
                    other.year,
                    other.month,
                    other.day,
                    other.hour,
                    other.minute,
                    other.second,
                    other.microsecond,
                )
            else:
                other = pendulum.instance(other)

        return self.diff(other, False) 

def _cmp(self, other, **kwargs):
        # Fix for pypy which compares using this method
        # which would lead to infinite recursion if we didn't override
        kwargs = {"tzinfo": self.tz}

        if _HAS_FOLD:
            kwargs["fold"] = self.fold

        dt = datetime.datetime(
            self.year,
            self.month,
            self.day,
            self.hour,
            self.minute,
            self.second,
            self.microsecond,
            **kwargs
        )

        return 0 if dt == other else 1 if dt > other else -1