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()