80 строки
2.8 KiB
Python
80 строки
2.8 KiB
Python
# -*- 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 unittest
|
|
from datetime import datetime
|
|
|
|
import mock
|
|
|
|
from airflow.utils import operator_helpers
|
|
|
|
|
|
class TestOperatorHelpers(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
super(TestOperatorHelpers, self).setUp()
|
|
self.dag_id = 'dag_id'
|
|
self.task_id = 'task_id'
|
|
self.execution_date = '2017-05-21T00:00:00'
|
|
self.dag_run_id = 'dag_run_id'
|
|
self.context = {
|
|
'dag_run': mock.MagicMock(
|
|
name='dag_run',
|
|
run_id=self.dag_run_id,
|
|
execution_date=datetime.strptime(self.execution_date,
|
|
'%Y-%m-%dT%H:%M:%S'),
|
|
),
|
|
'task_instance': mock.MagicMock(
|
|
name='task_instance',
|
|
task_id=self.task_id,
|
|
dag_id=self.dag_id,
|
|
execution_date=datetime.strptime(self.execution_date,
|
|
'%Y-%m-%dT%H:%M:%S'),
|
|
),
|
|
}
|
|
|
|
def test_context_to_airflow_vars_empty_context(self):
|
|
self.assertDictEqual(operator_helpers.context_to_airflow_vars({}), {})
|
|
|
|
def test_context_to_airflow_vars_all_context(self):
|
|
self.assertDictEqual(
|
|
operator_helpers.context_to_airflow_vars(self.context),
|
|
{
|
|
'airflow.ctx.dag_id': self.dag_id,
|
|
'airflow.ctx.execution_date': self.execution_date,
|
|
'airflow.ctx.task_id': self.task_id,
|
|
'airflow.ctx.dag_run_id': self.dag_run_id,
|
|
}
|
|
)
|
|
|
|
self.assertDictEqual(
|
|
operator_helpers.context_to_airflow_vars(self.context,
|
|
in_env_var_format=True),
|
|
{
|
|
'AIRFLOW_CTX_DAG_ID': self.dag_id,
|
|
'AIRFLOW_CTX_EXECUTION_DATE': self.execution_date,
|
|
'AIRFLOW_CTX_TASK_ID': self.task_id,
|
|
'AIRFLOW_CTX_DAG_RUN_ID': self.dag_run_id,
|
|
}
|
|
)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|