Ejemplo n.º 1
0
    def test_missing_arguments(
        self,
        missing_arg,
        input_data,
        voice,
        audio_config,
        target_bucket_name,
        target_filename,
        mock_text_to_speech_hook,
        mock_gcp_hook,
    ):
        with self.assertRaises(AirflowException) as e:
            CloudTextToSpeechSynthesizeOperator(
                project_id="project-id",
                input_data=input_data,
                voice=voice,
                audio_config=audio_config,
                target_bucket_name=target_bucket_name,
                target_filename=target_filename,
                task_id="id",
            ).execute(context={"task_instance": Mock()})

        err = e.exception
        self.assertIn(missing_arg, str(err))
        mock_text_to_speech_hook.assert_not_called()
        mock_gcp_hook.assert_not_called()
Ejemplo n.º 2
0
    def test_synthesize_text_green_path(self, mock_text_to_speech_hook,
                                        mock_gcp_hook):
        mocked_response = Mock()
        type(mocked_response).audio_content = PropertyMock(
            return_value=b"audio")

        mock_text_to_speech_hook.return_value.synthesize_speech.return_value = mocked_response
        mock_gcp_hook.return_value.upload.return_value = True

        CloudTextToSpeechSynthesizeOperator(
            project_id=PROJECT_ID,
            gcp_conn_id=GCP_CONN_ID,
            input_data=INPUT,
            voice=VOICE,
            audio_config=AUDIO_CONFIG,
            target_bucket_name=TARGET_BUCKET_NAME,
            target_filename=TARGET_FILENAME,
            task_id="id",
        ).execute(context={"task_instance": Mock()})

        mock_text_to_speech_hook.assert_called_once_with(
            gcp_conn_id="gcp-conn-id")
        mock_gcp_hook.assert_called_once_with(
            google_cloud_storage_conn_id="gcp-conn-id")
        mock_text_to_speech_hook.return_value.synthesize_speech.assert_called_once_with(
            input_data=INPUT,
            voice=VOICE,
            audio_config=AUDIO_CONFIG,
            retry=None,
            timeout=None)
        mock_gcp_hook.return_value.upload.assert_called_once_with(
            bucket_name=TARGET_BUCKET_NAME,
            object_name=TARGET_FILENAME,
            filename=ANY)
CONFIG = {"encoding": "LINEAR16", "language_code": "en_US"}
AUDIO = {"uri": "gs://{bucket}/{object}".format(bucket=BUCKET_NAME, object=FILENAME)}
# [END howto_operator_speech_to_text_api_arguments]

default_args = {"start_date": dates.days_ago(1)}

with models.DAG(
    "example_gcp_speech_to_text",
    default_args=default_args,
    schedule_interval=None,  # Override to match your needs
    tags=['example'],
) as dag:
    text_to_speech_synthesize_task = CloudTextToSpeechSynthesizeOperator(
        project_id=GCP_PROJECT_ID,
        input_data=INPUT,
        voice=VOICE,
        audio_config=AUDIO_CONFIG,
        target_bucket_name=BUCKET_NAME,
        target_filename=FILENAME,
        task_id="text_to_speech_synthesize_task",
    )
    # [START howto_operator_speech_to_text_recognize]
    speech_to_text_recognize_task2 = CloudSpeechToTextRecognizeSpeechOperator(
        config=CONFIG,
        audio=AUDIO,
        task_id="speech_to_text_recognize_task"
    )
    # [END howto_operator_speech_to_text_recognize]

    text_to_speech_synthesize_task >> speech_to_text_recognize_task2