[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

@@ -25,3 +25,4 @@ patch = mock.patch
 Mock = mock.Mock
 MagicMock = mock.MagicMock
 PropertyMock = mock.PropertyMock
+ANY = mock.ANY

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

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

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

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