class TestCloudTranslateHook(unittest.TestCase):
def setUp(self):
with mock.patch(
'airflow.gcp.hooks.translate.CloudTranslateHook.__init__',
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudTranslateHook(gcp_conn_id='test')
@mock.patch("airflow.gcp.hooks.translate.CloudTranslateHook.client_info",
new_callable=mock.PropertyMock)
@mock.patch(
"airflow.gcp.hooks.translate.CloudTranslateHook._get_credentials")
@mock.patch("airflow.gcp.hooks.translate.Client")
def test_translate_client_creation(self, mock_client, mock_get_creds,
mock_client_info):
result = self.hook.get_conn()
mock_client.assert_called_once_with(
credentials=mock_get_creds.return_value,
client_info=mock_client_info.return_value)
self.assertEqual(mock_client.return_value, result)
self.assertEqual(self.hook._client, result)
@mock.patch('airflow.gcp.hooks.translate.CloudTranslateHook.get_conn')
def test_translate_called(self, get_conn):
# Given
translate_method = get_conn.return_value.translate
translate_method.return_value = {
'translatedText': 'Yellowing self Gęśle',
'detectedSourceLanguage': 'pl',
'model': 'base',
'input': 'zażółć gęślą jaźń',
}
# When
result = self.hook.translate(
values=['zażółć gęślą jaźń'],
target_language='en',
format_='text',
source_language=None,
model='base',
)
# Then
self.assertEqual(
result,
{
'translatedText': 'Yellowing self Gęśle',
'detectedSourceLanguage': 'pl',
'model': 'base',
'input': 'zażółć gęślą jaźń',
},
)
translate_method.assert_called_once_with(
values=['zażółć gęślą jaźń'],
target_language='en',
format_='text',
source_language=None,
model='base',
)
def execute(self, context):
_hook = CloudTranslateHook(gcp_conn_id=self.gcp_conn_id)
try:
translation = _hook.translate(
values=self.values,
target_language=self.target_language,
format_=self.format_,
source_language=self.source_language,
model=self.model,
)
self.log.debug("Translation %s", translation)
return translation
except ValueError as e:
self.log.error('An error has been thrown from translate method:')
self.log.error(e)
raise AirflowException(e)
def execute(self, context):
speech_to_text_hook = CloudSpeechToTextHook(
gcp_conn_id=self.gcp_conn_id)
translate_hook = CloudTranslateHook(gcp_conn_id=self.gcp_conn_id)
recognize_result = speech_to_text_hook.recognize_speech(
config=self.config, audio=self.audio)
recognize_dict = MessageToDict(recognize_result)
self.log.info("Recognition operation finished")
if not recognize_dict['results']:
self.log.info("No recognition results")
return {}
self.log.debug("Recognition result: %s", recognize_dict)
try:
transcript = recognize_dict['results'][0]['alternatives'][0][
'transcript']
except KeyError as key:
raise AirflowException(
"Wrong response '{}' returned - it should contain {} field".
format(recognize_dict, key))
try:
translation = translate_hook.translate(
values=transcript,
target_language=self.target_language,
format_=self.format_,
source_language=self.source_language,
model=self.model)
self.log.info('Translated output: %s', translation)
return translation
except ValueError as e:
self.log.error(
'An error has been thrown from translate speech method:')
self.log.error(e)
raise AirflowException(e)
def setUp(self):
with mock.patch(
'airflow.gcp.hooks.translate.CloudTranslateHook.__init__',
new=mock_base_gcp_hook_default_project_id,
):
self.hook = CloudTranslateHook(gcp_conn_id='test')
