diff --git a/tests/providers/google/cloud/sensors/test_bigquery.py b/tests/providers/google/cloud/sensors/test_bigquery.py index c531a70e62a12..1ee469d62802d 100644 --- a/tests/providers/google/cloud/sensors/test_bigquery.py +++ b/tests/providers/google/cloud/sensors/test_bigquery.py @@ -16,7 +16,7 @@ # under the License. from __future__ import annotations -from unittest import TestCase, mock +from unittest import mock import pytest @@ -37,7 +37,7 @@ TEST_IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] -class TestBigqueryTableExistenceSensor(TestCase): +class TestBigqueryTableExistenceSensor: @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook") def test_passing_arguments_to_hook(self, mock_hook): task = BigQueryTableExistenceSensor( @@ -64,7 +64,7 @@ def test_passing_arguments_to_hook(self, mock_hook): ) -class TestBigqueryTablePartitionExistenceSensor(TestCase): +class TestBigqueryTablePartitionExistenceSensor: @mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook") def test_passing_arguments_to_hook(self, mock_hook): task = BigQueryTablePartitionExistenceSensor( @@ -104,7 +104,7 @@ def context(): yield context -class TestBigQueryTableExistenceAsyncSensor(TestCase): +class TestBigQueryTableExistenceAsyncSensor: def test_big_query_table_existence_sensor_async(self): """ Asserts that a task is deferred and a BigQueryTableExistenceTrigger will be fired diff --git a/tests/providers/google/cloud/sensors/test_bigquery_dts.py b/tests/providers/google/cloud/sensors/test_bigquery_dts.py index 270f15a7bf393..c4ea153001c4e 100644 --- a/tests/providers/google/cloud/sensors/test_bigquery_dts.py +++ b/tests/providers/google/cloud/sensors/test_bigquery_dts.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock from unittest.mock import MagicMock as MM @@ -35,7 +34,7 @@ GCP_CONN_ID = "google_cloud_default" -class TestBigQueryDataTransferServiceTransferRunSensor(unittest.TestCase): +class TestBigQueryDataTransferServiceTransferRunSensor: @mock.patch( "airflow.providers.google.cloud.sensors.bigquery_dts.BiqQueryDataTransferServiceHook", return_value=MM(get_transfer_run=MM(return_value=MM(state=TransferState.FAILED))), diff --git a/tests/providers/google/cloud/sensors/test_cloud_composer.py b/tests/providers/google/cloud/sensors/test_cloud_composer.py index c00521c1ebe85..8062da44b9b9a 100644 --- a/tests/providers/google/cloud/sensors/test_cloud_composer.py +++ b/tests/providers/google/cloud/sensors/test_cloud_composer.py @@ -17,7 +17,7 @@ from __future__ import annotations -from unittest import TestCase, mock +from unittest import mock import pytest @@ -30,7 +30,7 @@ TEST_REGION = "region" -class TestCloudComposerEnvironmentSensor(TestCase): +class TestCloudComposerEnvironmentSensor: def test_cloud_composer_existence_sensor_async(self): """ Asserts that a task is deferred and a CloudComposerExecutionTrigger will be fired diff --git a/tests/providers/google/cloud/sensors/test_cloud_storage_transfer_service.py b/tests/providers/google/cloud/sensors/test_cloud_storage_transfer_service.py index 86db78c3045e8..cbf348b324b67 100644 --- a/tests/providers/google/cloud/sensors/test_cloud_storage_transfer_service.py +++ b/tests/providers/google/cloud/sensors/test_cloud_storage_transfer_service.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock from parameterized import parameterized @@ -35,7 +34,7 @@ JOB_NAME = "job-name/123" -class TestGcpStorageTransferOperationWaitForJobStatusSensor(unittest.TestCase): +class TestGcpStorageTransferOperationWaitForJobStatusSensor: @mock.patch( "airflow.providers.google.cloud.sensors.cloud_storage_transfer_service.CloudDataTransferServiceHook" ) diff --git a/tests/providers/google/cloud/sensors/test_dataflow.py b/tests/providers/google/cloud/sensors/test_dataflow.py index 07043b9753ab3..54de120a30fe9 100644 --- a/tests/providers/google/cloud/sensors/test_dataflow.py +++ b/tests/providers/google/cloud/sensors/test_dataflow.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock import pytest @@ -41,7 +40,7 @@ TEST_IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] -class TestDataflowJobStatusSensor(unittest.TestCase): +class TestDataflowJobStatusSensor: @parameterized.expand( [ (DataflowJobStatus.JOB_STATE_DONE, DataflowJobStatus.JOB_STATE_DONE, True), @@ -107,7 +106,7 @@ def test_poke_raise_exception(self, mock_hook): ) -class TestDataflowJobMetricsSensor(unittest.TestCase): +class TestDataflowJobMetricsSensor: @parameterized.expand( [ (DataflowJobStatus.JOB_STATE_RUNNING, True), @@ -149,7 +148,7 @@ def test_poke(self, job_current_state, fail_on_terminal_state, mock_hook): callback.assert_called_once_with(mock_fetch_job_metrics_by_id.return_value.__getitem__.return_value) -class DataflowJobMessagesSensorTest(unittest.TestCase): +class DataflowJobMessagesSensorTest: @parameterized.expand( [ (DataflowJobStatus.JOB_STATE_RUNNING, True), @@ -225,7 +224,7 @@ def test_poke_raise_exception(self, mock_hook): callback.assert_not_called() -class DataflowJobAutoScalingEventsSensorTest(unittest.TestCase): +class DataflowJobAutoScalingEventsSensorTest: @parameterized.expand( [ (DataflowJobStatus.JOB_STATE_RUNNING, True), diff --git a/tests/providers/google/cloud/sensors/test_datafusion.py b/tests/providers/google/cloud/sensors/test_datafusion.py index 195f02068131d..16ebb0d71f9b6 100644 --- a/tests/providers/google/cloud/sensors/test_datafusion.py +++ b/tests/providers/google/cloud/sensors/test_datafusion.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock import pytest @@ -39,7 +38,7 @@ FAILURE_STATUSES = {"FAILED"} -class TestCloudDataFusionPipelineStateSensor(unittest.TestCase): +class TestCloudDataFusionPipelineStateSensor: @parameterized.expand( [ (PipelineStates.COMPLETED, PipelineStates.COMPLETED, True), diff --git a/tests/providers/google/cloud/sensors/test_dataplex.py b/tests/providers/google/cloud/sensors/test_dataplex.py index dc409bc43a092..9871a2f4c5d39 100644 --- a/tests/providers/google/cloud/sensors/test_dataplex.py +++ b/tests/providers/google/cloud/sensors/test_dataplex.py @@ -16,7 +16,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock import pytest @@ -40,7 +39,7 @@ IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"] -class TestDataplexTaskStateSensor(unittest.TestCase): +class TestDataplexTaskStateSensor: def create_task(self, state: int): task = mock.Mock() task.state = state diff --git a/tests/providers/google/cloud/sensors/test_dataproc.py b/tests/providers/google/cloud/sensors/test_dataproc.py index d05d9f62b587c..61d6d55c919ba 100644 --- a/tests/providers/google/cloud/sensors/test_dataproc.py +++ b/tests/providers/google/cloud/sensors/test_dataproc.py @@ -16,7 +16,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock from unittest.mock import Mock @@ -39,7 +38,7 @@ TIMEOUT = 120 -class TestDataprocJobSensor(unittest.TestCase): +class TestDataprocJobSensor: def create_job(self, state: int): job = mock.Mock() job.status = mock.Mock() @@ -186,7 +185,7 @@ def test_wait_timeout_raise_exception(self, mock_hook): sensor.poke(context={}) -class TestDataprocBatchSensor(unittest.TestCase): +class TestDataprocBatchSensor: def create_batch(self, state: int): batch = mock.Mock() batch.state = mock.Mock() diff --git a/tests/providers/google/cloud/sensors/test_gcs.py b/tests/providers/google/cloud/sensors/test_gcs.py index bf586dfe1c02f..642461f63a635 100644 --- a/tests/providers/google/cloud/sensors/test_gcs.py +++ b/tests/providers/google/cloud/sensors/test_gcs.py @@ -18,7 +18,7 @@ from __future__ import annotations from datetime import datetime, timedelta, timezone -from unittest import TestCase, mock +from unittest import mock import pendulum import pytest @@ -75,7 +75,7 @@ def next_time_side_effect(): mock_time = mock.Mock(side_effect=next_time_side_effect) -class TestGoogleCloudStorageObjectSensor(TestCase): +class TestGoogleCloudStorageObjectSensor: @mock.patch("airflow.providers.google.cloud.sensors.gcs.GCSHook") def test_should_pass_argument_to_hook(self, mock_hook): task = GCSObjectExistenceSensor( @@ -99,7 +99,7 @@ def test_should_pass_argument_to_hook(self, mock_hook): mock_hook.return_value.exists.assert_called_once_with(TEST_BUCKET, TEST_OBJECT, DEFAULT_RETRY) -class TestGoogleCloudStorageObjectSensorAsync(TestCase): +class TestGoogleCloudStorageObjectSensorAsync: def test_gcs_object_existence_sensor_async(self): """ Asserts that a task is deferred and a GCSBlobTrigger will be fired @@ -139,7 +139,7 @@ def test_gcs_object_existence_sensor_async_execute_complete(self): mock_log_info.assert_called_with("File %s was found in bucket %s.", TEST_OBJECT, TEST_BUCKET) -class TestTsFunction(TestCase): +class TestTsFunction: def test_should_support_datetime(self): context = { "dag": DAG(dag_id=TEST_DAG_ID, schedule=timedelta(days=5)), @@ -159,7 +159,7 @@ def test_should_support_cron(self): assert pendulum.instance(datetime(2019, 2, 24)).isoformat() == result.isoformat() -class TestGoogleCloudStorageObjectUpdatedSensor(TestCase): +class TestGoogleCloudStorageObjectUpdatedSensor: @mock.patch("airflow.providers.google.cloud.sensors.gcs.GCSHook") def test_should_pass_argument_to_hook(self, mock_hook): task = GCSObjectUpdateSensor( @@ -182,7 +182,7 @@ def test_should_pass_argument_to_hook(self, mock_hook): assert result is True -class TestGoogleCloudStoragePrefixSensor(TestCase): +class TestGoogleCloudStoragePrefixSensor: @mock.patch("airflow.providers.google.cloud.sensors.gcs.GCSHook") def test_should_pass_arguments_to_hook(self, mock_hook): task = GCSObjectsWithPrefixExistenceSensor( @@ -253,8 +253,8 @@ def test_execute_timeout(self, mock_hook): mock_hook.return_value.list.assert_called_once_with(TEST_BUCKET, prefix=TEST_PREFIX) -class TestGCSUploadSessionCompleteSensor(TestCase): - def setUp(self): +class TestGCSUploadSessionCompleteSensor: + def setup_method(self): self.dag = DAG( TEST_DAG_ID + "test_schedule_dag_once", schedule="@once", diff --git a/tests/providers/google/cloud/sensors/test_looker.py b/tests/providers/google/cloud/sensors/test_looker.py index ceaf3c8ea98d2..567340f0864d6 100644 --- a/tests/providers/google/cloud/sensors/test_looker.py +++ b/tests/providers/google/cloud/sensors/test_looker.py @@ -16,7 +16,6 @@ # under the License. from __future__ import annotations -import unittest from unittest import mock import pytest @@ -33,7 +32,7 @@ TEST_JOB_ID = "123" -class TestLookerCheckPdtBuildSensor(unittest.TestCase): +class TestLookerCheckPdtBuildSensor: @mock.patch(SENSOR_PATH.format("LookerHook")) def test_done(self, mock_hook): mock_hook.return_value.pdt_build_status.return_value = {"status": JobStatus.DONE.value} diff --git a/tests/providers/google/cloud/sensors/test_pubsub.py b/tests/providers/google/cloud/sensors/test_pubsub.py index cc4191e0ee6c2..952758578c3e7 100644 --- a/tests/providers/google/cloud/sensors/test_pubsub.py +++ b/tests/providers/google/cloud/sensors/test_pubsub.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from typing import Any from unittest import mock @@ -32,7 +31,7 @@ TEST_SUBSCRIPTION = "test-subscription" -class TestPubSubPullSensor(unittest.TestCase): +class TestPubSubPullSensor: def _generate_messages(self, count): return [ ReceivedMessage( diff --git a/tests/providers/google/cloud/sensors/test_tasks.py b/tests/providers/google/cloud/sensors/test_tasks.py index a3a08e26cdc31..59356ce9e6127 100644 --- a/tests/providers/google/cloud/sensors/test_tasks.py +++ b/tests/providers/google/cloud/sensors/test_tasks.py @@ -17,7 +17,6 @@ # under the License. from __future__ import annotations -import unittest from typing import Any from unittest import mock @@ -35,7 +34,7 @@ FULL_TASK_PATH = "projects/test-project/locations/asia-east2/queues/test-queue/tasks/test-task" -class TestCloudTasksEmptySensor(unittest.TestCase): +class TestCloudTasksEmptySensor: @mock.patch("airflow.providers.google.cloud.sensors.tasks.CloudTasksHook") def test_queue_empty(self, mock_hook):