def test_missing_file(self, mock_exists):
with self.assertRaisesRegex(
AirflowException,
re.escape(
"File a.json was not found. Check the configuration of your Secrets backend."
),
):
local_filesystem.load_connections_dict("a.json")
def _import_helper(file_path):
"""Helps import connections from a file"""
connections_dict = load_connections_dict(file_path)
with create_session() as session:
for conn_id, conn_values in connections_dict.items():
if session.query(Connection).filter(
Connection.conn_id == conn_id).first():
print(
f'Could not import connection {conn_id}: connection already exists.'
)
continue
allowed_fields = [
'extra',
'description',
'conn_id',
'login',
'conn_type',
'host',
'password',
'schema',
'port',
'uri',
'extra_dejson',
]
filtered_connection_values = {
key: value
for key, value in conn_values.items() if key in allowed_fields
}
connection = _create_connection(conn_id,
filtered_connection_values)
session.add(connection)
session.commit()
print(f'Imported connection {conn_id}')
def test_yaml_file_should_load_connection_extras(self, file_content, expected_extras):
with mock_local_file(file_content):
connections_by_conn_id = local_filesystem.load_connections_dict("a.yaml")
connection_uris_by_conn_id = {
conn_id: connection.extra_dejson for conn_id, connection in connections_by_conn_id.items()
}
self.assertEqual(expected_extras, connection_uris_by_conn_id)
def test_json_file_should_load_connection(self, file_content, expected_connection_uris):
with mock_local_file(json.dumps(file_content)):
connections_by_conn_id = local_filesystem.load_connections_dict("a.json")
connection_uris_by_conn_id = {
conn_id: connection.get_uri()
for conn_id, connection in connections_by_conn_id.items()
}
self.assertEqual(expected_connection_uris, connection_uris_by_conn_id)
def test_yaml_file_should_load_connection(self, file_content,
expected_connection_uris):
with mock_local_file(file_content):
connections_by_conn_id = local_filesystem.load_connections_dict(
"a.yaml")
connection_uris_by_conn_id = {
conn_id: connection.get_uri()
for conn_id, connection in connections_by_conn_id.items()
}
assert expected_connection_uris == connection_uris_by_conn_id
def test_yaml_file_should_load_connection(self, file_content,
expected_attrs_dict):
with mock_local_file(file_content):
connections_by_conn_id = local_filesystem.load_connections_dict(
"a.yaml")
for conn_id, connection in connections_by_conn_id.items():
expected_attrs = expected_attrs_dict[conn_id]
actual_attrs = {
k: getattr(connection, k)
for k in expected_attrs.keys()
}
assert actual_attrs == expected_attrs
def _import_helper(file_path):
"""Load connections from a file and save them to the DB. On collision, skip."""
connections_dict = load_connections_dict(file_path)
with create_session() as session:
for conn_id, conn in connections_dict.items():
if session.query(Connection).filter(
Connection.conn_id == conn_id).first():
print(
f'Could not import connection {conn_id}: connection already exists.'
)
continue
session.add(conn)
session.commit()
print(f'Imported connection {conn_id}')
def test_ensure_unique_connection_yaml(self, file_content):
with mock_local_file(file_content):
with self.assertRaises(ConnectionNotUnique):
local_filesystem.load_connections_dict("a.yaml")
def test_ensure_unique_connection_json(self, file_content):
with mock_local_file(json.dumps(file_content)):
with self.assertRaises(ConnectionNotUnique):
local_filesystem.load_connections_dict("a.json")
def test_yaml_invalid_extra(self, file_content, expected_message):
with mock_local_file(file_content):
with self.assertRaisesRegex(AirflowException,
re.escape(expected_message)):
local_filesystem.load_connections_dict("a.yaml")
def test_env_file_invalid_input(self, file_content,
expected_connection_uris):
with mock_local_file(json.dumps(file_content)):
with self.assertRaisesRegex(AirflowException,
re.escape(expected_connection_uris)):
local_filesystem.load_connections_dict("a.json")
def test_env_file_invalid_format(self, content, expected_message):
with mock_local_file(content):
with self.assertRaisesRegex(AirflowFileParseException,
re.escape(expected_message)):
local_filesystem.load_connections_dict("a.env")
def test_ensure_unique_connection_env(self, file_content):
with mock_local_file(file_content):
with pytest.raises(ConnectionNotUnique):
local_filesystem.load_connections_dict("a.env")
def test_yaml_invalid_extra(self, file_content, expected_message):
with mock_local_file(file_content):
with pytest.raises(AirflowException,
match=re.escape(expected_message)):
local_filesystem.load_connections_dict("a.yaml")
def test_env_file_invalid_input(self, file_content,
expected_connection_uris):
with mock_local_file(json.dumps(file_content)):
with pytest.raises(AirflowException,
match=re.escape(expected_connection_uris)):
local_filesystem.load_connections_dict("a.json")
