def test_get_variable_first_try(self, mock_env_get, mock_meta_get):
"""
Test if Variable is present in Environment Variable, it does not look for it in
Metastore DB
"""
mock_env_get.return_value = [["something"]] # returns nonempty list
Variable.get_variable_from_secrets("fake_var_key")
mock_env_get.assert_called_once_with(key="fake_var_key")
mock_meta_get.not_called()
def test_get_variable_second_try(self, mock_env_get, mock_meta_get):
"""
Test if Variable is not present in Environment Variable, it then looks for it in
Metastore DB
"""
mock_env_get.return_value = None
Variable.get_variable_from_secrets("fake_var_key")
mock_meta_get.assert_called_once_with(key="fake_var_key")
mock_env_get.assert_called_once_with(key="fake_var_key")
def variables_import(args):
"""Imports variables from a given file"""
try:
vars_map = load_variables(args.file)
except AirflowException as e:
print(e)
sys.exit(1)
if not args.conflict_disposition:
disposition = DIS_RESTRICT
else:
disposition = args.conflict_disposition
var_status_map = {DIS_OVERWRITE: [], DIS_IGNORE: [], CREATED: []}
try:
with create_session() as session:
for key, val in vars_map.items():
serialize = isinstance(val, (list, dict))
vars_row = Variable.get_variable_from_secrets(key)
if not vars_row:
Variable.set(key=key,
value=val,
session=session,
serialize_json=serialize)
session.flush()
var_status_map[CREATED].append(key)
elif disposition == DIS_OVERWRITE:
Variable.set(key=key,
value=val,
session=session,
serialize_json=serialize)
session.flush()
var_status_map[DIS_OVERWRITE].append(key)
elif disposition == DIS_IGNORE:
var_status_map[DIS_IGNORE].append(key)
else:
msg = "\nVariable with `key`={key} already exists"
msg = msg.format(key=key)
raise AirflowException(msg)
print(_prep_import_status_msgs(var_status_map))
except (SQLAlchemyError, AirflowException) as e:
print(e)
session.rollback()
finally:
session.close()
def test_variables_import_conflict_disposition_overwrite(
self, file_content, expected_value):
"""Test variables_import command with --conflict-disposition overwrite"""
with mock_local_file(file_content[0], 'a.json'):
variable_command.variables_import(
self.parser.parse_args(['variables', 'import', 'a.json']))
with mock_local_file(file_content[1], 'a.json'):
variable_command.variables_import(
self.parser.parse_args([
'variables', 'import', 'a.json', '--conflict-disposition',
'overwrite'
]))
var = Variable.get_variable_from_secrets(expected_value["key"])
self.assertEqual(var, expected_value["val"])
