Skip to content

Migrate tests in google/cloud/sensors from unittest to pytest #29642

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dimberman merged 1 commit into from
Feb 27, 2023
Merged

Migrate tests in google/cloud/sensors from unittest to pytest #29642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions tests/providers/google/cloud/sensors/test_bigquery.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
# under the License.
from __future__ import annotations

from unittest import TestCase, mock
from unittest import mock

import pytest

Expand All @@ -37,7 +37,7 @@
TEST_IMPERSONATION_CHAIN = ["ACCOUNT_1", "ACCOUNT_2", "ACCOUNT_3"]


class TestBigqueryTableExistenceSensor(TestCase):
class TestBigqueryTableExistenceSensor:

@dimberman dimberman Feb 20, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If we're switching from unittest to pytest then why are we even using test classes at all? We might as well go all the way and make the tests into functions. I'd rather make this change now while we're already modifying the code.

@Taragolis Taragolis Feb 20, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I hope that classes vs functions it is not new tabs vs spaces 🤣
Personally I do not have any preferences, outside of Airflow I use both of them depends on situation.

If module contain a lot of different functions / classes which need to be tested then classes very nice to use as namespace so developers could easily found tests for specific part, like:

class TestAwesomeHook:

    def test_init(self, ...):
        ...

    def test_init_failure_reason_1(self, ...):
        ...

    def test_get_conn(self, ...):
        ...

class TestAnotherAwesomeHook:

    def test_init(self, ...):
        ...

    def test_init_failure_reason_1(self, ...):
        ...

    def test_get_conn(self, ...):
        ...

vs

def test_awesome_hook_init(...):
    ...

def test_awesome_hook_init_failure_reason_1(...):
    ...

def test_another_awesome_hook_init(...):
    ...

def test_another_awesome_hook_get_conn(...):
   ...

def test_awesome_hook_init_failure_reason_1(...):
    ...

def test_awesome_hook_get_conn(...):
   ...

@pankajastro pankajastro Feb 21, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

especially in the case of providers, I feel that class-based tests help because we have many operators/sensors in a single file so if tests are class-based it improves the readability of tests and adding/removing/updating/searching tests became a bit easy. otherwise, If I'll make small changes in the existing operator most of the time goes in figuring out do I need to add new tests or can update existing and where shall I add the new tests in a test file

@vedantlodha vedantlodha Feb 25, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I agree! Tried retaking a look, and in my opinion as well, having class names for sensors makes the tests lot more readable. @dimberman I would request you to re-evaluate your request to remove classes. Thanks

@dimberman dimberman Feb 27, 2023

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ahhh ok that makes sense. I'm always a fan of minimizing nesting as much as possible but I think this is reasonable considering there are multiple classes in the file

@mock.patch("airflow.providers.google.cloud.sensors.bigquery.BigQueryHook")
def test_passing_arguments_to_hook(self, mock_hook):
task = BigQueryTableExistenceSensor(
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_bigquery_dts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock
from unittest.mock import MagicMock as MM

Expand All @@ -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))),
Expand Down
4 changes: 2 additions & 2 deletions tests/providers/google/cloud/sensors/test_cloud_composer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@

from __future__ import annotations

from unittest import TestCase, mock
from unittest import mock

import pytest

Expand All @@ -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
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_cloud_storage_transfer_service.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock

from parameterized import parameterized
Expand All @@ -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"
)
Expand Down
9 changes: 4 additions & 5 deletions tests/providers/google/cloud/sensors/test_dataflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock

import pytest
Expand All @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down Expand Up @@ -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),
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_datafusion.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock

import pytest
Expand All @@ -39,7 +38,7 @@
FAILURE_STATUSES = {"FAILED"}


class TestCloudDataFusionPipelineStateSensor(unittest.TestCase):
class TestCloudDataFusionPipelineStateSensor:
@parameterized.expand(
[
(PipelineStates.COMPLETED, PipelineStates.COMPLETED, True),
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_dataplex.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock

import pytest
Expand All @@ -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
Expand Down
5 changes: 2 additions & 3 deletions tests/providers/google/cloud/sensors/test_dataproc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock
from unittest.mock import Mock

Expand All @@ -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()
Expand Down Expand Up @@ -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()
Expand Down
16 changes: 8 additions & 8 deletions tests/providers/google/cloud/sensors/test_gcs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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)),
Expand All @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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",
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_looker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
# under the License.
from __future__ import annotations

import unittest
from unittest import mock

import pytest
Expand All @@ -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}
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_pubsub.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from typing import Any
from unittest import mock

Expand All @@ -32,7 +31,7 @@
TEST_SUBSCRIPTION = "test-subscription"


class TestPubSubPullSensor(unittest.TestCase):
class TestPubSubPullSensor:
def _generate_messages(self, count):
return [
ReceivedMessage(
Expand Down
3 changes: 1 addition & 2 deletions tests/providers/google/cloud/sensors/test_tasks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
# under the License.
from __future__ import annotations

import unittest
from typing import Any
from unittest import mock

Expand All @@ -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):

Expand Down