Migrate Google Cloud hooks test to pytest

tests/providers/google/cloud/hooks/test_automl.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 from google.api_core.gapic_v1.method import DEFAULT
@@ -51,8 +50,8 @@
 MASK = {"field": "mask"}
 
 
-class TestAuoMLHook(unittest.TestCase):
-    def setUp(self) -> None:
+class TestAuoMLHook:
+    def setup_method(self):
         with mock.patch(
             "airflow.providers.google.cloud.hooks.automl.GoogleBaseHook.__init__",
             new=mock_base_gcp_hook_no_default_project_id,

tests/providers/google/cloud/hooks/test_bigquery.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import re
-import unittest
 from datetime import datetime
 from unittest import mock
 
@@ -27,7 +26,6 @@
 from google.cloud.bigquery import DEFAULT_RETRY, DatasetReference, Table, TableReference
 from google.cloud.bigquery.dataset import AccessEntry, Dataset, DatasetListItem
 from google.cloud.exceptions import NotFound
-from parameterized import parameterized
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery import (
@@ -256,24 +254,25 @@ def test_run_query_without_sql_fails(self, _):
         ):
             self.hook.run_query(sql=None)
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "schema_update_options, write_disposition",
         [
             (["ALLOW_FIELD_ADDITION"], "WRITE_APPEND"),
             (["ALLOW_FIELD_RELAXATION"], "WRITE_APPEND"),
             (["ALLOW_FIELD_ADDITION", "ALLOW_FIELD_RELAXATION"], "WRITE_APPEND"),
             (["ALLOW_FIELD_ADDITION"], "WRITE_TRUNCATE"),
             (["ALLOW_FIELD_RELAXATION"], "WRITE_TRUNCATE"),
             (["ALLOW_FIELD_ADDITION", "ALLOW_FIELD_RELAXATION"], "WRITE_TRUNCATE"),
-        ]
+        ],
     )
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")
     def test_run_query_schema_update_options(
         self,
+        mock_insert,
+        _,
         schema_update_options,
         write_disposition,
-        mock_insert,
-        mock_get_service,
     ):
         self.hook.run_query(
             sql="query",
@@ -285,7 +284,8 @@ def test_run_query_schema_update_options(
         assert kwargs["configuration"]["query"]["schemaUpdateOptions"] == schema_update_options
         assert kwargs["configuration"]["query"]["writeDisposition"] == write_disposition
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "schema_update_options, write_disposition, expected_regex",
         [
             (
                 ["INCORRECT_OPTION"],
@@ -307,15 +307,11 @@ def test_run_query_schema_update_options(
                 r"schema_update_options is only allowed if write_disposition is "
                 r"'WRITE_APPEND' or 'WRITE_TRUNCATE'",
             ),
-        ]
+        ],
     )
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     def test_run_query_schema_update_options_incorrect(
-        self,
-        schema_update_options,
-        write_disposition,
-        expected_regex,
-        mock_get_service,
+        self, _, schema_update_options, write_disposition, expected_regex
     ):
         with pytest.raises(ValueError, match=expected_regex):
             self.hook.run_query(
@@ -325,15 +321,10 @@ def test_run_query_schema_update_options_incorrect(
                 write_disposition=write_disposition,
             )
 
-    @parameterized.expand([(True,), (False,)])
+    @pytest.mark.parametrize("bool_val", [True, False])
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.get_service")
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")
-    def test_api_resource_configs(
-        self,
-        bool_val,
-        mock_insert,
-        _,
-    ):
+    def test_api_resource_configs(self, mock_insert, _, bool_val):
         self.hook.run_query("query", api_resource_configs={"query": {"useQueryCache": bool_val}})
         _, kwargs = mock_insert.call_args
         assert kwargs["configuration"]["query"]["useQueryCache"] is bool_val
@@ -397,9 +388,9 @@ def test_validate_src_fmt_configs(self):
             "compatibility_val" in src_fmt_configs
         ), "_validate_src_fmt_configs should add backward_compatibility config"
 
-    @parameterized.expand([("AVRO",), ("PARQUET",), ("NEWLINE_DELIMITED_JSON",), ("DATASTORE_BACKUP",)])
+    @pytest.mark.parametrize("fmt", ["AVRO", "PARQUET", "NEWLINE_DELIMITED_JSON", "DATASTORE_BACKUP"])
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")
-    def test_run_load_with_non_csv_as_src_fmt(self, fmt, _):
+    def test_run_load_with_non_csv_as_src_fmt(self, _, fmt):
 
         try:
             self.hook.run_load(
@@ -408,8 +399,8 @@ def test_run_load_with_non_csv_as_src_fmt(self, fmt, _):
                 source_format=fmt,
                 autodetect=True,
             )
-        except ValueError:
-            self.fail("run_load() raised ValueError unexpectedly!")
+        except ValueError as ex:
+            pytest.fail("run_load() raised ValueError unexpectedly!", ex)
 
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.insert_job")
     def test_run_extract(self, mock_insert):
@@ -949,19 +940,20 @@ def test_job_id_validity(self, mock_md5, test_dag_id, expected_job_id):
         assert job_id == expected_job_id
 
 
-class TestBigQueryTableSplitter(unittest.TestCase):
+class TestBigQueryTableSplitter:
     def test_internal_need_default_project(self):
         with pytest.raises(Exception, match="INTERNAL: No default project is specified"):
             split_tablename("dataset.table", None)
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "project_expected, dataset_expected, table_expected, table_input",
         [
             ("project", "dataset", "table", "dataset.table"),
             ("alternative", "dataset", "table", "alternative:dataset.table"),
             ("alternative", "dataset", "table", "alternative.dataset.table"),
             ("alt1:alt", "dataset", "table", "alt1:alt.dataset.table"),
             ("alt1:alt", "dataset", "table", "alt1:alt:dataset.table"),
-        ]
+        ],
     )
     def test_split_tablename(self, project_expected, dataset_expected, table_expected, table_input):
         default_project_id = "project"
@@ -970,7 +962,8 @@ def test_split_tablename(self, project_expected, dataset_expected, table_expecte
         assert dataset_expected == dataset
         assert table_expected == table
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "table_input, var_name, exception_message",
         [
             ("alt1:alt2:alt3:dataset.table", None, "Use either : or . to specify project got {}"),
             (
@@ -994,7 +987,7 @@ def test_split_tablename(self, project_expected, dataset_expected, table_expecte
                 r"Format exception for var_x: Expect format of "
                 r"\(<project\.\|<project:\)<dataset>.<table>, got {}",
             ),
-        ]
+        ],
     )
     def test_invalid_syntax(self, table_input, var_name, exception_message):
         default_project_id = "project"
@@ -1972,37 +1965,38 @@ def test_run_load_with_kms(self, mock_insert):
         )
 
 
-class TestBigQueryBaseCursorMethodsDeprecationWarning(unittest.TestCase):
-    @parameterized.expand(
+class TestBigQueryBaseCursorMethodsDeprecationWarning:
+    @pytest.mark.parametrize(
+        "func_name",
         [
-            ("create_empty_table",),
-            ("create_empty_dataset",),
-            ("get_dataset_tables",),
-            ("delete_dataset",),
-            ("create_external_table",),
-            ("patch_table",),
-            ("insert_all",),
-            ("update_dataset",),
-            ("patch_dataset",),
-            ("get_dataset_tables_list",),
-            ("get_datasets_list",),
-            ("get_dataset",),
-            ("run_grant_dataset_view_access",),
-            ("run_table_upsert",),
-            ("run_table_delete",),
-            ("get_tabledata",),
-            ("get_schema",),
-            ("poll_job_complete",),
-            ("cancel_query",),
-            ("run_with_configuration",),
-            ("run_load",),
-            ("run_copy",),
-            ("run_extract",),
-            ("run_query",),
-        ]
+            "create_empty_table",
+            "create_empty_dataset",
+            "get_dataset_tables",
+            "delete_dataset",
+            "create_external_table",
+            "patch_table",
+            "insert_all",
+            "update_dataset",
+            "patch_dataset",
+            "get_dataset_tables_list",
+            "get_datasets_list",
+            "get_dataset",
+            "run_grant_dataset_view_access",
+            "run_table_upsert",
+            "run_table_delete",
+            "get_tabledata",
+            "get_schema",
+            "poll_job_complete",
+            "cancel_query",
+            "run_with_configuration",
+            "run_load",
+            "run_copy",
+            "run_extract",
+            "run_query",
+        ],
     )
     @mock.patch("airflow.providers.google.cloud.hooks.bigquery.BigQueryHook")
-    def test_deprecation_warning(self, func_name, mock_bq_hook):
+    def test_deprecation_warning(self, mock_bq_hook, func_name):
         args, kwargs = [1], {"param1": "val1"}
         new_path = re.escape(f"`airflow.providers.google.cloud.hooks.bigquery.BigQueryHook.{func_name}`")
         message_pattern = rf"This method is deprecated\.\s+Please use {new_path}"

tests/providers/google/cloud/hooks/test_bigquery_dts.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from copy import deepcopy
 from unittest import mock
 
@@ -50,8 +49,8 @@
 RUN_ID = "id1234"
 
 
-class BigQueryDataTransferHookTestCase(unittest.TestCase):
-    def setUp(self) -> None:
+class TestBigQueryDataTransferHook:
+    def setup_method(self):
         with mock.patch(
             "airflow.providers.google.cloud.hooks.bigquery_dts.GoogleBaseHook.__init__",
             new=mock_base_gcp_hook_no_default_project_id,

tests/providers/google/cloud/hooks/test_bigtable.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 from unittest.mock import PropertyMock
 
@@ -50,8 +49,8 @@
 ]
 
 
-class TestBigtableHookNoDefaultProjectId(unittest.TestCase):
-    def setUp(self):
+class TestBigtableHookNoDefaultProjectId:
+    def setup_method(self):
         with mock.patch(
             "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
             new=mock_base_gcp_hook_no_default_project_id,
@@ -146,8 +145,8 @@ def test_delete_table_overridden_project_id(self, get_client):
         table_delete_method.assert_called_once_with()
 
 
-class TestBigtableHookDefaultProjectId(unittest.TestCase):
-    def setUp(self):
+class TestBigtableHookDefaultProjectId:
+    def setup_method(self):
         with mock.patch(
             "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
             new=mock_base_gcp_hook_default_project_id,
@@ -322,15 +321,13 @@ def test_create_instance_with_one_replica_cluster_production(
         )
         cluster.assert_has_calls(
             [
-                unittest.mock.call(
+                mock.call(
                     cluster_id=CBT_CLUSTER,
                     location_id=CBT_ZONE,
                     serve_nodes=1,
                     default_storage_type=enums.StorageType.SSD,
                 ),
-                unittest.mock.call(
-                    CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD
-                ),
+                mock.call(CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD),
             ],
             any_order=True,
         )
@@ -365,12 +362,10 @@ def test_create_instance_with_one_replica_cluster_development(
         )
         cluster.assert_has_calls(
             [
-                unittest.mock.call(
+                mock.call(
                     cluster_id=CBT_CLUSTER, location_id=CBT_ZONE, default_storage_type=enums.StorageType.SSD
                 ),
-                unittest.mock.call(
-                    CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD
-                ),
+                mock.call(CBT_REPLICA_CLUSTER_ID, CBT_REPLICA_CLUSTER_ZONE, 1, enums.StorageType.SSD),
             ],
             any_order=True,
         )
@@ -404,15 +399,15 @@ def test_create_instance_with_multiple_replica_clusters(
         )
         cluster.assert_has_calls(
             [
-                unittest.mock.call(
+                mock.call(
                     cluster_id=CBT_CLUSTER,
                     location_id=CBT_ZONE,
                     serve_nodes=1,
                     default_storage_type=enums.StorageType.SSD,
                 ),
-                unittest.mock.call("replica-1", "us-west1-a", 1, enums.StorageType.SSD),
-                unittest.mock.call("replica-2", "us-central1-f", 1, enums.StorageType.SSD),
-                unittest.mock.call("replica-3", "us-east1-d", 1, enums.StorageType.SSD),
+                mock.call("replica-1", "us-west1-a", 1, enums.StorageType.SSD),
+                mock.call("replica-2", "us-central1-f", 1, enums.StorageType.SSD),
+                mock.call("replica-3", "us-east1-d", 1, enums.StorageType.SSD),
             ],
             any_order=True,
         )

tests/providers/google/cloud/hooks/test_cloud_build.py

@@ -21,7 +21,6 @@
 """
 from __future__ import annotations
 
-import unittest
 from concurrent.futures import Future
 from unittest import mock
 
@@ -60,8 +59,8 @@
 TRIGGER_ID = "32488e7f-09d6-4fe9-a5fb-4ca1419a6e7a"
 
 
-class TestCloudBuildHook(unittest.TestCase):
-    def setUp(self):
+class TestCloudBuildHook:
+    def setup_method(self):
         with mock.patch(
             "airflow.providers.google.common.hooks.base_google.GoogleBaseHook.__init__",
             new=mock_base_gcp_hook_no_default_project_id,

tests/providers/google/cloud/hooks/test_cloud_composer.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 import pytest
@@ -60,8 +59,8 @@ def mock_init(*args, **kwargs):
     pass
 
 
-class TestCloudComposerHook(unittest.TestCase):
-    def setUp(self) -> None:
+class TestCloudComposerHook:
+    def setup_method(self):
         with mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_init):
             self.hook = CloudComposerHook(gcp_conn_id="test")
 
@@ -197,7 +196,7 @@ def test_list_image_versions(self, mock_client) -> None:
 
 
 class TestCloudComposerAsyncHook:
-    def setup_method(self, method) -> None:
+    def setup_method(self, method):
         with async_mock.patch(BASE_STRING.format("GoogleBaseHook.__init__"), new=mock_init):
             self.hook = CloudComposerAsyncHook(gcp_conn_id="test")