def test_get_config_make_data_dir(self, mock_get, mock_file, mock_os_makedirs):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_ini_missing_section)
configuration.get_config(config_filename='config.ini',
make_data_dir=False)
mock_os_makedirs.assert_not_called()
def test_get_config_yaml_format(self, mock_get, mock_file, mock_os_makedirs):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_yaml)
conf = configuration.get_config(config_filename='config.yaml')
mock_file.assert_called_with('/usr/local/airflow/config/utils/config.yaml', 'r')
self.assertEqual(
mock_os_makedirs.mock_calls,
[call('/usr/local/airflow/data/utils', exist_ok=True),
call('/usr/local/airflow/data/utils/dir_dev1', exist_ok=True),
call('/usr/local/airflow/data/utils/dir_dev2', exist_ok=True)])
self.assertEqual(
conf,
{'config_path': '/usr/local/airflow/config/utils',
'dag_email': '*****@*****.**',
'dag_email_on_failure': False,
'dag_email_on_retry': False,
'data_dirs_extra': {'test_path1': 'dir_dev1', 'test_path_dev2': 'dir_dev2'},
'data_path': '/usr/local/airflow/data/utils',
'email_path': '/usr/local/airflow/emails/utils',
'environment': 'DEV',
'sql_path': '/usr/local/airflow/sql/utils',
'test_path1': '/usr/local/airflow/data/utils/dir_dev1',
'test_path_dev2': '/usr/local/airflow/data/utils/dir_dev2',
'var1': 'value_dev1',
'var_dev2': 'value_dev2',
'var_list_dev': ['item1_dev', 'item2_dev']})
def test_get_config_ini_format(self, mock_get, mock_file, mock_os_makedirs):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_ini)
conf = configuration.get_config(config_filename='config.ini')
mock_file.assert_called_with('/usr/local/airflow/config/utils/config.ini', encoding=None)
self.assertIsInstance(conf, configparser.SectionProxy)
self.assertEqual(
mock_os_makedirs.mock_calls,
[call('/usr/local/airflow/data/utils', exist_ok=True),
call('/usr/local/airflow/data/utils/dir_dev1', exist_ok=True),
call('/usr/local/airflow/data/utils/dir_dev2', exist_ok=True)])
self.assertEqual(
dict(conf),
{'config_path': '/usr/local/airflow/config/utils',
'dag_email': '*****@*****.**',
'dag_email_on_failure': 'False',
'dag_email_on_retry': 'False',
'data_dirs_extra': '{"test_path1": "dir_dev1", "test_path_dev2": "dir_dev2"}',
'data_path': '/usr/local/airflow/data/utils',
'email_path': '/usr/local/airflow/emails/utils',
'environment': 'DEV',
'sql_path': '/usr/local/airflow/sql/utils',
'test_path1': '/usr/local/airflow/data/utils/dir_dev1',
'test_path_dev2': '/usr/local/airflow/data/utils/dir_dev2',
'var1': 'value_dev1',
'var2': 'value2',
'var_dev2': 'value_dev2'})
def test_get_config_project_directory(self, mock_get, mock_file, _):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_ini_missing_section)
conf = configuration.get_config(config_filename='config.ini', project_directory='project_dir')
self.assertEqual(
dict(conf),
{'config_path': '/usr/local/airflow/config/project_dir',
'dag_email': '*****@*****.**',
'dag_email_on_failure': 'False',
'dag_email_on_retry': 'False',
'data_path': '/usr/local/airflow/data/project_dir',
'email_path': '/usr/local/airflow/emails/project_dir',
'environment': 'DEV',
'sql_path': '/usr/local/airflow/sql/project_dir',
'var1': 'value1'})
def test_get_config_data_dirs_extra(self, mock_get, mock_file, _):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_ini_missing_section)
conf = configuration.get_config(config_filename='config.ini',
data_dirs_extra={'test_path1': 'dir_dev1',
'test_path_dev2': 'dir_dev2'})
self.assertEqual(
dict(conf),
{'config_path': '/usr/local/airflow/config/utils',
'dag_email': '*****@*****.**',
'dag_email_on_failure': 'False',
'dag_email_on_retry': 'False',
'data_path': '/usr/local/airflow/data/utils',
'email_path': '/usr/local/airflow/emails/utils',
'environment': 'DEV',
'sql_path': '/usr/local/airflow/sql/utils',
'test_path1': '/usr/local/airflow/data/utils/dir_dev1',
'test_path_dev2': '/usr/local/airflow/data/utils/dir_dev2',
'var1': 'value1'})
polling_period_seconds=60 * 3,
notebook_params={"config_path": file_spec['config_path'],
"file_date": date_str,
"file_path": "{}/{}".format(
input_s3_dir, str(file_spec.get('unzipped_search_expr',
file_spec['search_expr'])).replace(".*", "*"))}
)
save_to_s3 >> import_file
download >> check_for_files >> unzip_files >> decrypt >> save_to_s3 >> remove_tmp_files
return dag
today_date = datetime.now()
# Get configuration settings for a project
base_conf = get_config(config_filename='config.yaml')
for config_filename in base_conf.get('config_filenames'):
config = copy.deepcopy(base_conf)
conf = get_config(config_filename=config_filename)
dict_merge(config, conf)
dag_id = config.get('dag_id').format(config_filename=config_filename.split('.')[0])
globals()[dag_id] = create_dag(
dag_id=dag_id,
description='DAG to download files from {} and load to database'.format(config_filename.split('.')[0]),
conf=config,
date=today_date)
def test_get_config_default_config_file(self, mock_get, mock_file, _):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_ini_missing_section)
configuration.get_config()
mock_file.assert_called_with('/usr/local/airflow/config/utils/config.ini', encoding=None)
def test_get_config_yaml_format_missing_section(self, mock_get, mock_file, _):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_yaml_missing_section)
conf = configuration.get_config(config_filename='config.yaml')
self.assertEqual(dict(conf), self.expected_missing_section_yaml)
def test_get_config_yml_extension(self, mock_get, mock_file, _):
mock_get.side_effect = self.config_vars.get
mock_file.return_value = io.StringIO(self.config_yaml)
configuration.get_config(config_filename='config.yml')
mock_file.assert_called_with('/usr/local/airflow/config/utils/config.yml', 'r')