From 36ff0d03bd0fd7f7d6711955df6d6b69c0cb617d Mon Sep 17 00:00:00 2001 From: Felix Uellendall Date: Tue, 2 Apr 2019 21:32:38 +0200 Subject: [PATCH] [AIRFLOW-4223] Fix mypy issue in gcp speech_to_text and text_to_speech hooks and operators (#5027) - adds mock.ANY to compat.py --- tests/compat.py | 1 + .../hooks/test_gcp_speech_to_text_hook.py | 13 ++------- .../hooks/test_gcp_text_to_speech_hook.py | 13 ++------- .../test_gcp_speech_to_text_operator.py | 22 +++++--------- .../test_gcp_text_to_speech_operator.py | 29 +++++++------------ 5 files changed, 24 insertions(+), 54 deletions(-) diff --git a/tests/compat.py b/tests/compat.py index 8269a0f92a..7fe554b355 100644 --- a/tests/compat.py +++ b/tests/compat.py @@ -25,3 +25,4 @@ patch = mock.patch Mock = mock.Mock MagicMock = mock.MagicMock PropertyMock = mock.PropertyMock +ANY = mock.ANY diff --git a/tests/contrib/hooks/test_gcp_speech_to_text_hook.py b/tests/contrib/hooks/test_gcp_speech_to_text_hook.py index 5de4684016..baecd0664d 100644 --- a/tests/contrib/hooks/test_gcp_speech_to_text_hook.py +++ b/tests/contrib/hooks/test_gcp_speech_to_text_hook.py @@ -21,16 +21,9 @@ import unittest from airflow.contrib.hooks.gcp_speech_to_text_hook import GCPSpeechToTextHook +from tests.compat import patch from tests.contrib.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id -try: - from unittest import mock -except ImportError: # pragma: no cover - try: - import mock - except ImportError: - mock = None - PROJECT_ID = "project-id" CONFIG = {"ecryption": "LINEAR16"} AUDIO = {"uri": "gs://bucket/object"} @@ -38,13 +31,13 @@ AUDIO = {"uri": "gs://bucket/object"} class TestTextToSpeechOperator(unittest.TestCase): def setUp(self): - with mock.patch( + with patch( "airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.gcp_speech_to_text_hook = GCPSpeechToTextHook(gcp_conn_id="test") - @mock.patch("airflow.contrib.hooks.gcp_speech_to_text_hook.GCPSpeechToTextHook.get_conn") + @patch("airflow.contrib.hooks.gcp_speech_to_text_hook.GCPSpeechToTextHook.get_conn") def test_synthesize_speech(self, get_conn): recognize_method = get_conn.return_value.recognize recognize_method.return_value = None diff --git a/tests/contrib/hooks/test_gcp_text_to_speech_hook.py b/tests/contrib/hooks/test_gcp_text_to_speech_hook.py index c217dbb8a8..eb4b5b75e7 100644 --- a/tests/contrib/hooks/test_gcp_text_to_speech_hook.py +++ b/tests/contrib/hooks/test_gcp_text_to_speech_hook.py @@ -21,16 +21,9 @@ import unittest from airflow.contrib.hooks.gcp_text_to_speech_hook import GCPTextToSpeechHook +from tests.compat import patch from tests.contrib.utils.base_gcp_mock import mock_base_gcp_hook_default_project_id -try: - from unittest import mock -except ImportError: # pragma: no cover - try: - import mock - except ImportError: - mock = None - INPUT = {"text": "test text"} VOICE = {"language_code": "en-US", "ssml_gender": "FEMALE"} AUDIO_CONFIG = {"audio_encoding": "MP3"} @@ -38,13 +31,13 @@ AUDIO_CONFIG = {"audio_encoding": "MP3"} class TestTextToSpeechHook(unittest.TestCase): def setUp(self): - with mock.patch( + with patch( "airflow.contrib.hooks.gcp_api_base_hook.GoogleCloudBaseHook.__init__", new=mock_base_gcp_hook_default_project_id, ): self.gcp_text_to_speech_hook = GCPTextToSpeechHook(gcp_conn_id="test") - @mock.patch("airflow.contrib.hooks.gcp_text_to_speech_hook.GCPTextToSpeechHook.get_conn") + @patch("airflow.contrib.hooks.gcp_text_to_speech_hook.GCPTextToSpeechHook.get_conn") def test_synthesize_speech(self, get_conn): synthesize_method = get_conn.return_value.synthesize_speech synthesize_method.return_value = None diff --git a/tests/contrib/operators/test_gcp_speech_to_text_operator.py b/tests/contrib/operators/test_gcp_speech_to_text_operator.py index c457158e71..0fdf479c89 100644 --- a/tests/contrib/operators/test_gcp_speech_to_text_operator.py +++ b/tests/contrib/operators/test_gcp_speech_to_text_operator.py @@ -21,15 +21,7 @@ import unittest from airflow import AirflowException from airflow.contrib.operators.gcp_speech_to_text_operator import GcpSpeechToTextRecognizeSpeechOperator - - -try: - from unittest import mock -except ImportError: - try: - import mock - except ImportError: - mock = None +from tests.compat import Mock, patch PROJECT_ID = "project-id" GCP_CONN_ID = "gcp-conn-id" @@ -38,40 +30,40 @@ AUDIO = {"uri": "gs://bucket/object"} class CloudSqlTest(unittest.TestCase): - @mock.patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") + @patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") def test_recognize_speech_green_path(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True GcpSpeechToTextRecognizeSpeechOperator( project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, config=CONFIG, audio=AUDIO, task_id="id" - ).execute(context={"task_instance": mock.Mock()}) + ).execute(context={"task_instance": Mock()}) mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID) mock_hook.return_value.recognize_speech.assert_called_once_with( config=CONFIG, audio=AUDIO, retry=None, timeout=None ) - @mock.patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") + @patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") def test_missing_config(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True with self.assertRaises(AirflowException) as e: GcpSpeechToTextRecognizeSpeechOperator( project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, audio=AUDIO, task_id="id" - ).execute(context={"task_instance": mock.Mock()}) + ).execute(context={"task_instance": Mock()}) err = e.exception self.assertIn("config", str(err)) mock_hook.assert_not_called() - @mock.patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") + @patch("airflow.contrib.operators.gcp_speech_to_text_operator.GCPSpeechToTextHook") def test_missing_audio(self, mock_hook): mock_hook.return_value.recognize_speech.return_value = True with self.assertRaises(AirflowException) as e: GcpSpeechToTextRecognizeSpeechOperator( project_id=PROJECT_ID, gcp_conn_id=GCP_CONN_ID, config=CONFIG, task_id="id" - ).execute(context={"task_instance": mock.Mock()}) + ).execute(context={"task_instance": Mock()}) err = e.exception self.assertIn("audio", str(err)) diff --git a/tests/contrib/operators/test_gcp_text_to_speech_operator.py b/tests/contrib/operators/test_gcp_text_to_speech_operator.py index bd27409417..d1f9b18bcb 100644 --- a/tests/contrib/operators/test_gcp_text_to_speech_operator.py +++ b/tests/contrib/operators/test_gcp_text_to_speech_operator.py @@ -22,17 +22,8 @@ import unittest from parameterized import parameterized from airflow import AirflowException - from airflow.contrib.operators.gcp_text_to_speech_operator import GcpTextToSpeechSynthesizeOperator - - -try: - from unittest import mock -except ImportError: - try: - import mock - except ImportError: - mock = None +from tests.compat import PropertyMock, Mock, patch, ANY PROJECT_ID = "project-id" GCP_CONN_ID = "gcp-conn-id" @@ -44,11 +35,11 @@ TARGET_FILENAME = "target_filename" class GcpTextToSpeechTest(unittest.TestCase): - @mock.patch("airflow.contrib.operators.gcp_text_to_speech_operator.GoogleCloudStorageHook") - @mock.patch("airflow.contrib.operators.gcp_text_to_speech_operator.GCPTextToSpeechHook") + @patch("airflow.contrib.operators.gcp_text_to_speech_operator.GoogleCloudStorageHook") + @patch("airflow.contrib.operators.gcp_text_to_speech_operator.GCPTextToSpeechHook") def test_synthesize_text_green_path(self, mock_text_to_speech_hook, mock_gcp_hook): - mocked_response = mock.Mock() - type(mocked_response).audio_content = mock.PropertyMock(return_value=b"audio") + 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 @@ -62,7 +53,7 @@ class GcpTextToSpeechTest(unittest.TestCase): target_bucket_name=TARGET_BUCKET_NAME, target_filename=TARGET_FILENAME, task_id="id", - ).execute(context={"task_instance": mock.Mock()}) + ).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") @@ -70,7 +61,7 @@ class GcpTextToSpeechTest(unittest.TestCase): input_data=INPUT, voice=VOICE, audio_config=AUDIO_CONFIG, retry=None, timeout=None ) mock_gcp_hook.return_value.upload.assert_called_once_with( - bucket=TARGET_BUCKET_NAME, object=TARGET_FILENAME, filename=mock.ANY + bucket=TARGET_BUCKET_NAME, object=TARGET_FILENAME, filename=ANY ) @parameterized.expand( @@ -82,8 +73,8 @@ class GcpTextToSpeechTest(unittest.TestCase): ("target_filename", INPUT, VOICE, AUDIO_CONFIG, TARGET_BUCKET_NAME, ""), ] ) - @mock.patch("airflow.contrib.operators.gcp_text_to_speech_operator.GoogleCloudStorageHook") - @mock.patch("airflow.contrib.operators.gcp_text_to_speech_operator.GCPTextToSpeechHook") + @patch("airflow.contrib.operators.gcp_text_to_speech_operator.GoogleCloudStorageHook") + @patch("airflow.contrib.operators.gcp_text_to_speech_operator.GCPTextToSpeechHook") def test_missing_arguments( self, missing_arg, @@ -104,7 +95,7 @@ class GcpTextToSpeechTest(unittest.TestCase): target_bucket_name=target_bucket_name, target_filename=target_filename, task_id="id", - ).execute(context={"task_instance": mock.Mock()}) + ).execute(context={"task_instance": Mock()}) err = e.exception self.assertIn(missing_arg, str(err))