def test_cli_delete_connections(self, session=None):
merge_conn(
Connection(
conn_id="new1",
conn_type="mysql",
description="mysql description",
host="mysql",
login="******",
password="",
schema="airflow",
),
session=session,
)
# Delete connections
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_delete(self.parser.parse_args(["connections", "delete", "new1"]))
stdout = stdout.getvalue()
# Check deletion stdout
self.assertIn("\tSuccessfully deleted `conn_id`=new1", stdout)
# Check deletions
result = session.query(Connection).filter(Connection.conn_id == "new1").first()
self.assertTrue(result is None)
def test_cli_delete_invalid_connection(self):
# Attempt to delete a non-existing connection
with pytest.raises(
SystemExit,
match=r"Did not find a connection with `conn_id`=fake"):
connection_command.connections_delete(
self.parser.parse_args(["connections", "delete", "fake"]))
def test_cli_delete_invalid_connection(self):
# Attempt to delete a non-existing connection
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_delete(self.parser.parse_args(["connections", "delete", "fake"]))
stdout = stdout.getvalue()
# Check deletion attempt stdout
self.assertIn("\tDid not find a connection with `conn_id`=fake", stdout)
def test_cli_connections_add_delete(self):
# TODO: We should not delete the entire database, but only reset the contents of the Connection table.
db.resetdb()
# Add connections:
uri = 'postgresql://*****:*****@host:5432/airflow'
with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
connection_command.connections_add(
self.parser.parse_args(
['connections', 'add', 'new1',
'--conn_uri=%s' % uri]))
connection_command.connections_add(
self.parser.parse_args(
['connections', 'add', 'new2',
'--conn_uri=%s' % uri]))
connection_command.connections_add(
self.parser.parse_args([
'connections', 'add', 'new3',
'--conn_uri=%s' % uri, '--conn_extra', "{'extra': 'yes'}"
]))
connection_command.connections_add(
self.parser.parse_args([
'connections', 'add', 'new4',
'--conn_uri=%s' % uri, '--conn_extra', "{'extra': 'yes'}"
]))
connection_command.connections_add(
self.parser.parse_args([
'connections', 'add', 'new5', '--conn_type=hive_metastore',
'--conn_login=airflow', '--conn_password=airflow',
'--conn_host=host', '--conn_port=9083',
'--conn_schema=airflow'
]))
connection_command.connections_add(
self.parser.parse_args([
'connections', 'add', 'new6', '--conn_uri', "",
'--conn_type=google_cloud_platform', '--conn_extra',
"{'extra': 'yes'}"
]))
stdout = mock_stdout.getvalue()
# Check addition stdout
lines = [l for l in stdout.split('\n') if len(l) > 0]
self.assertListEqual(
lines, [("\tSuccessfully added `conn_id`=new1 : " +
"postgresql://*****:*****@host:5432/airflow"),
("\tSuccessfully added `conn_id`=new2 : " +
"postgresql://*****:*****@host:5432/airflow"),
("\tSuccessfully added `conn_id`=new3 : " +
"postgresql://*****:*****@host:5432/airflow"),
("\tSuccessfully added `conn_id`=new4 : " +
"postgresql://*****:*****@host:5432/airflow"),
("\tSuccessfully added `conn_id`=new5 : " +
"hive_metastore://airflow:airflow@host:9083/airflow"),
("\tSuccessfully added `conn_id`=new6 : " +
"google_cloud_platform://:@:")])
# Attempt to add duplicate
with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
connection_command.connections_add(
self.parser.parse_args(
['connections', 'add', 'new1',
'--conn_uri=%s' % uri]))
stdout = mock_stdout.getvalue()
# Check stdout for addition attempt
lines = [l for l in stdout.split('\n') if len(l) > 0]
self.assertListEqual(lines, [
"\tA connection with `conn_id`=new1 already exists",
])
# Attempt to add without providing conn_uri
with self.assertRaises(SystemExit) as exc:
connection_command.connections_add(
self.parser.parse_args(['connections', 'add', 'new']))
self.assertEqual(
exc.exception.code,
"The following args are required to add a connection: ['conn_uri or conn_type']"
)
# Prepare to add connections
session = settings.Session()
extra = {
'new1': None,
'new2': None,
'new3': "{'extra': 'yes'}",
'new4': "{'extra': 'yes'}"
}
# Add connections
for index in range(1, 6):
conn_id = 'new%s' % index
result = (session.query(Connection).filter(
Connection.conn_id == conn_id).first())
result = (result.conn_id, result.conn_type, result.host,
result.port, result.get_extra())
if conn_id in ['new1', 'new2', 'new3', 'new4']:
self.assertEqual(
result,
(conn_id, 'postgres', 'host', 5432, extra[conn_id]))
elif conn_id == 'new5':
self.assertEqual(
result, (conn_id, 'hive_metastore', 'host', 9083, None))
elif conn_id == 'new6':
self.assertEqual(result, (conn_id, 'google_cloud_platform',
None, None, "{'extra': 'yes'}"))
# Delete connections
with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new1']))
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new2']))
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new3']))
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new4']))
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new5']))
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'new6']))
stdout = mock_stdout.getvalue()
# Check deletion stdout
lines = [l for l in stdout.split('\n') if len(l) > 0]
self.assertListEqual(lines, [
"\tSuccessfully deleted `conn_id`=new1",
"\tSuccessfully deleted `conn_id`=new2",
"\tSuccessfully deleted `conn_id`=new3",
"\tSuccessfully deleted `conn_id`=new4",
"\tSuccessfully deleted `conn_id`=new5",
"\tSuccessfully deleted `conn_id`=new6"
])
# Check deletions
for index in range(1, 7):
conn_id = 'new%s' % index
result = (session.query(Connection).filter(
Connection.conn_id == conn_id).first())
self.assertTrue(result is None)
# Attempt to delete a non-existing connection
with mock.patch('sys.stdout', new_callable=io.StringIO) as mock_stdout:
connection_command.connections_delete(
self.parser.parse_args(['connections', 'delete', 'fake']))
stdout = mock_stdout.getvalue()
# Check deletion attempt stdout
lines = [l for l in stdout.split('\n') if len(l) > 0]
self.assertListEqual(lines, [
"\tDid not find a connection with `conn_id`=fake",
])
session.close()
