Fix mypy tests providers part 2 (#20111)

apache · Dec 7, 2021 · 139be53 · 139be53
1 parent 50bf536
commit 139be53
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 10 deletions.
diff --git a/tests/providers/cncf/kubernetes/utils/test_pod_launcher.py b/tests/providers/cncf/kubernetes/utils/test_pod_launcher.py
@@ -20,7 +20,7 @@
 import pendulum
 import pytest
 from kubernetes.client.rest import ApiException
-from requests.exceptions import BaseHTTPError
+from urllib3.exceptions import HTTPError as BaseHTTPError
 
 from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.utils.pod_launcher import PodLauncher, PodStatus

diff --git a/tests/providers/google/cloud/operators/test_datacatalog.py b/tests/providers/google/cloud/operators/test_datacatalog.py
@@ -610,7 +610,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
 class TestCloudDataCatalogListTagsOperator(TestCase):
     @mock.patch(
         "airflow.providers.google.cloud.operators.datacatalog.CloudDataCatalogHook",
-        **{"return_value.list_tags.return_value": [TEST_TAG]},  # type: ignore
+        return_value=mock.MagicMock(list_tags=mock.MagicMock(return_value=[TEST_TAG])),
     )
     def test_assert_valid_hook_call(self, mock_hook) -> None:
         task = CloudDataCatalogListTagsOperator(
@@ -777,7 +777,7 @@ class TestCloudDataCatalogUpdateTagOperator(TestCase):
     def test_assert_valid_hook_call(self, mock_hook) -> None:
         task = CloudDataCatalogUpdateTagOperator(
             task_id="task_id",
-            tag=TEST_TAG_ID,
+            tag=Tag(name=TEST_TAG_ID),
             update_mask=TEST_UPDATE_MASK,
             location=TEST_LOCATION,
             entry_group=TEST_ENTRY_GROUP_ID,
@@ -796,7 +796,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
             impersonation_chain=TEST_IMPERSONATION_CHAIN,
         )
         mock_hook.return_value.update_tag.assert_called_once_with(
-            tag=TEST_TAG_ID,
+            tag=Tag(name=TEST_TAG_ID),
             update_mask=TEST_UPDATE_MASK,
             location=TEST_LOCATION,
             entry_group=TEST_ENTRY_GROUP_ID,
@@ -814,7 +814,7 @@ class TestCloudDataCatalogUpdateTagTemplateOperator(TestCase):
     def test_assert_valid_hook_call(self, mock_hook) -> None:
         task = CloudDataCatalogUpdateTagTemplateOperator(
             task_id="task_id",
-            tag_template=TEST_TAG_TEMPLATE_ID,
+            tag_template=TagTemplate(name=TEST_TAG_TEMPLATE_ID),
             update_mask=TEST_UPDATE_MASK,
             location=TEST_LOCATION,
             tag_template_id=TEST_TAG_TEMPLATE_ID,
@@ -831,7 +831,7 @@ def test_assert_valid_hook_call(self, mock_hook) -> None:
             impersonation_chain=TEST_IMPERSONATION_CHAIN,
         )
         mock_hook.return_value.update_tag_template.assert_called_once_with(
-            tag_template=TEST_TAG_TEMPLATE_ID,
+            tag_template=TagTemplate(name=TEST_TAG_TEMPLATE_ID),
             update_mask=TEST_UPDATE_MASK,
             location=TEST_LOCATION,
             tag_template_id=TEST_TAG_TEMPLATE_ID,

diff --git a/tests/providers/google/cloud/operators/test_dataproc_metastore.py b/tests/providers/google/cloud/operators/test_dataproc_metastore.py
@@ -110,7 +110,7 @@ def test_assert_valid_hook_call(self, mock_metadata_import, mock_hook) -> None:
             project_id=GCP_PROJECT_ID,
             region=GCP_LOCATION,
             service_id=TEST_SERVICE_ID,
-            metadata_import=TEST_METADATA_IMPORT,
+            metadata_import=mock_metadata_import(name=TEST_METADATA_IMPORT),
             metadata_import_id=TEST_METADATA_IMPORT_ID,
             retry=TEST_RETRY,
             timeout=TEST_TIMEOUT,
@@ -126,7 +126,7 @@ def test_assert_valid_hook_call(self, mock_metadata_import, mock_hook) -> None:
             project_id=GCP_PROJECT_ID,
             region=GCP_LOCATION,
             service_id=TEST_SERVICE_ID,
-            metadata_import=TEST_METADATA_IMPORT,
+            metadata_import=mock_metadata_import(name=TEST_METADATA_IMPORT),
             metadata_import_id=TEST_METADATA_IMPORT_ID,
             request_id=None,
             retry=TEST_RETRY,

diff --git a/tests/providers/google/cloud/sensors/test_bigquery_dts.py b/tests/providers/google/cloud/sensors/test_bigquery_dts.py
@@ -18,6 +18,7 @@
 
 import unittest
 from unittest import mock
+from unittest.mock import MagicMock as MM
 
 from google.cloud.bigquery_datatransfer_v1 import TransferState
 
@@ -31,7 +32,7 @@
 class TestBigQueryDataTransferServiceTransferRunSensor(unittest.TestCase):
     @mock.patch(
         "airflow.providers.google.cloud.sensors.bigquery_dts.BiqQueryDataTransferServiceHook",
-        **{'return_value.get_transfer_run.return_value.state': TransferState.FAILED},
+        return_value=MM(get_transfer_run=MM(return_value=MM(state=TransferState.FAILED))),
     )
     def test_poke_returns_false(self, mock_hook):
         op = BigQueryDataTransferServiceTransferRunSensor(
@@ -55,7 +56,7 @@ def test_poke_returns_false(self, mock_hook):
 
     @mock.patch(
         "airflow.providers.google.cloud.sensors.bigquery_dts.BiqQueryDataTransferServiceHook",
-        **{'return_value.get_transfer_run.return_value.state': TransferState.SUCCEEDED},
+        return_value=MM(get_transfer_run=MM(return_value=MM(state=TransferState.SUCCEEDED))),
     )
     def test_poke_returns_true(self, mock_hook):
         op = BigQueryDataTransferServiceTransferRunSensor(