def test_cli_connections_add_duplicate(self):
# Attempt to add duplicate
connection_command.connections_add(
self.parser.parse_args(
["connections", "add", "new1",
"--conn_uri=%s" % TEST_URL]))
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_add(
self.parser.parse_args(
["connections", "add", "new1",
"--conn_uri=%s" % TEST_URL]))
stdout = stdout.getvalue()
# Check stdout for addition attempt
self.assertIn("\tA connection with `conn_id`=new1 already exists",
stdout)
def test_cli_connection_add(self, cmd, expected_output, expected_conn):
with redirect_stdout(io.StringIO()) as stdout:
connection_command.connections_add(self.parser.parse_args(cmd))
stdout = stdout.getvalue()
self.assertIn(expected_output, stdout)
conn_id = cmd[2]
with create_session() as session:
comparable_attrs = [
"conn_type",
"host",
"is_encrypted",
"is_extra_encrypted",
"login",
"port",
"schema",
]
current_conn = session.query(Connection).filter(Connection.conn_id == conn_id).first()
self.assertEqual(expected_conn, {attr: getattr(current_conn, attr) for attr in comparable_attrs})
def test_cli_connections_add_invalid_uri(self):
# Attempt to add with invalid uri
with self.assertRaisesRegex(SystemExit, r"The URI provided to --conn-uri is invalid: nonsense_uri"):
connection_command.connections_add(
self.parser.parse_args(["connections", "add", "new1", "--conn-uri=%s" % "nonsense_uri"])
)
def test_cli_connections_add_delete_with_missing_parameters(self):
# Attempt to add without providing conn_uri
with self.assertRaisesRegex(
SystemExit, r"The following args are required to add a connection: \['conn-uri or conn-type'\]"
):
connection_command.connections_add(self.parser.parse_args(["connections", "add", "new1"]))
def test_cli_connections_add_invalid_uri(self):
# Attempt to add with invalid uri
with pytest.raises(SystemExit, match=r"The URI provided to --conn-uri is invalid: nonsense_uri"):
connection_command.connections_add(
self.parser.parse_args(["connections", "add", "new1", f"--conn-uri={'nonsense_uri'}"])
)
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()
