def test_connection_type_not_supported(self, mock_get_connection):
conn = get_airflow_connection("NOT_SUPPORT")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
with self.assertRaises(AirflowConfigException):
hook.get_conn()
def test_custom_connection_with_no_connection_func(self,
mock_get_connection):
conn = get_airflow_connection("CUSTOM")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
with self.assertRaises(AirflowConfigException):
hook.get_conn()
def test_custom_connection(self, mock_get_connection):
conn = get_airflow_connection("CUSTOM")
mock_get_connection.return_value = conn
mocked_channel = self.channel_mock.return_value
hook = GrpcHook("grpc_default",
custom_connection_func=self.custom_conn_func)
channel = hook.get_conn()
self.assertEqual(channel, mocked_channel)
def test_connection_with_port(self, mock_get_connection,
mock_insecure_channel):
conn = get_airflow_connection_with_port()
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_insecure_channel.return_value = mocked_channel
channel = hook.get_conn()
expected_url = "test.com:1234"
mock_insecure_channel.assert_called_once_with(expected_url)
self.assertEqual(channel, mocked_channel)
def test_no_auth_connection(self, mock_get_connection,
mock_insecure_channel):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_insecure_channel.return_value = mocked_channel
channel = hook.get_conn()
expected_url = "test:8080"
mock_insecure_channel.assert_called_once_with(expected_url)
assert channel == mocked_channel
def test_connection_with_interceptors(self, mock_insecure_channel,
mock_get_connection,
mock_intercept_channel):
conn = get_airflow_connection()
mock_get_connection.return_value = conn
mocked_channel = self.channel_mock.return_value
hook = GrpcHook("grpc_default", interceptors=["test1"])
mock_insecure_channel.return_value = mocked_channel
mock_intercept_channel.return_value = mocked_channel
channel = hook.get_conn()
self.assertEqual(channel, mocked_channel)
mock_intercept_channel.assert_called_once_with(mocked_channel, "test1")
def test_connection_with_jwt(self, mock_secure_channel,
mock_google_default_auth, mock_google_cred,
mock_get_connection):
conn = get_airflow_connection(auth_type="JWT_GOOGLE")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_google_default_auth.return_value = (mock_credential_object, "")
mock_google_cred.return_value = mock_credential_object
channel = hook.get_conn()
expected_url = "test:8080"
mock_google_cred.assert_called_once_with(mock_credential_object)
mock_secure_channel.assert_called_once_with(mock_credential_object,
None, expected_url)
self.assertEqual(channel, mocked_channel)
def test_connection_with_tls(self, mock_secure_channel,
mock_channel_credentials, mock_get_connection,
mock_open):
conn = get_airflow_connection(auth_type="TLS",
credential_pem_file="pem")
mock_get_connection.return_value = conn
mock_open.return_value = StringIO('credential')
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_channel_credentials.return_value = mock_credential_object
channel = hook.get_conn()
expected_url = "test:8080"
mock_open.assert_called_once_with("pem")
mock_channel_credentials.assert_called_once_with('credential')
mock_secure_channel.assert_called_once_with(expected_url,
mock_credential_object)
self.assertEqual(channel, mocked_channel)
def test_connection_with_google_oauth(self, mock_secure_channel,
mock_google_default_auth,
mock_google_auth_request,
mock_get_connection):
conn = get_airflow_connection(auth_type="OATH_GOOGLE",
scopes="grpc,gcs")
mock_get_connection.return_value = conn
hook = GrpcHook("grpc_default")
mocked_channel = self.channel_mock.return_value
mock_secure_channel.return_value = mocked_channel
mock_credential_object = "test_credential_object"
mock_google_default_auth.return_value = (mock_credential_object, "")
mock_google_auth_request.return_value = "request"
channel = hook.get_conn()
expected_url = "test:8080"
mock_google_default_auth.assert_called_once_with(
scopes=["grpc", "gcs"])
mock_secure_channel.assert_called_once_with(mock_credential_object,
"request", expected_url)
self.assertEqual(channel, mocked_channel)