Replace `unittests` in amazon provider tests by pure `pytest` #27970

Taragolis · 2022-11-28T19:51:43Z

Migrate Amazon provider's tests to pytest.

All changes are more or less straightforward:

Get rid of unittests.TestCase class and TestCase.assert* methods
Replace decorator parameterized.expand by pytest.mark.parametrize. I
replace requests_mock.mock decorator by requests_mock fixture
Convert TestCase.subTest to parametrize tests
Convert classes setUp* and tearDown* to appropriate pytest methods

See additional findings, info about significant changes and potential follow up in comments

o-nikolas

It's very difficult to review such a large PR for correctness, but I spot checked a handful of code and it seems like it looks good

Taragolis

As @o-nikolas mention that this kind of PR very hard to check (without losing concentration).

I add comments for anything which can not be classified as 'straightforward changes'

Taragolis · 2022-11-29T04:21:53Z

tests/providers/amazon/aws/operators/test_eks.py


-class TestEksCreateClusterOperator(unittest.TestCase):
-    def setUp(self) -> None:
+class TestEksCreateClusterOperator:


replace TestCase.subTest by parametrize tests

Rename method nodegroup_setUp to nodegroup_setup

Rename method fargate_profile_setup to fargate_profile_setup

Taragolis · 2022-11-29T04:23:24Z

tests/providers/amazon/aws/operators/test_eks.py

-class TestEksCreateFargateProfileOperator(unittest.TestCase):
-    def setUp(self) -> None:
-        self.create_fargate_profile_params: CreateFargateProfileParams = dict(  # type: ignore
+class TestEksCreateFargateProfileOperator:


replace TestCase.subTest by parametrize tests

Taragolis · 2022-11-29T04:24:09Z

tests/providers/amazon/aws/operators/test_eks.py

+        mock_create_fargate_profile.assert_called_with(**convert_keys(parameters))
+
+
+class TestEksCreateNodegroupOperator:


replace TestCase.subTest by parametrize tests

Taragolis · 2022-11-29T04:25:45Z

tests/providers/amazon/aws/operators/test_eks.py

-
-        mock_list_nodegroups.assert_called_once


This actually do nothing right now if convert to mock_list_nodegroups.assert_called_once() than it failed.
Remove this for now

Taragolis · 2022-11-29T04:27:22Z

tests/providers/amazon/aws/operators/test_emr_add_steps.py

+    @pytest.mark.parametrize(
+        "job_flow_id, job_flow_name",
+        [
+            pytest.param("j-8989898989", "test_cluster", id="both-specified"),
+            pytest.param(None, None, id="both-none"),
+        ],
+    )
+    def test_validate_mutually_exclusive_args(self, job_flow_id, job_flow_name):
+        error_message = r"Exactly one of job_flow_id or job_flow_name must be specified\."
+        with pytest.raises(AirflowException, match=error_message):
+            EmrAddStepsOperator(
+                task_id="test_validate_mutually_exclusive_args",
+                job_flow_id=job_flow_id,
+                job_flow_name=job_flow_name,
+            )
+


Add this test from #27858

Taragolis · 2022-11-29T04:37:29Z

tests/providers/amazon/aws/operators/test_sagemaker_transform.py

+        self.create_transform_params = copy.deepcopy(CREATE_TRANSFORM_PARAMS)
+        self.create_model_params = copy.deepcopy(CREATE_MODEL_PARAMS)
+        self.config = {"Model": self.create_model_params, "Transform": self.create_transform_params}


After migrate to pytest some test failed due to mutability of test parameters so this parameters recreate for each test case.

Taragolis · 2022-11-29T04:38:55Z

tests/providers/amazon/aws/sensors/test_athena.py

-    @mock.patch.object(AthenaHook, "poll_query_status", side_effect=("SUCCEEDED",))
-    def test_poke_success(self, mock_poll_query_status):
-        assert self.sensor.poke({})
-
-    @mock.patch.object(AthenaHook, "poll_query_status", side_effect=("RUNNING",))
-    def test_poke_running(self, mock_poll_query_status):
-        assert not self.sensor.poke({})
-
-    @mock.patch.object(AthenaHook, "poll_query_status", side_effect=("QUEUED",))
-    def test_poke_queued(self, mock_poll_query_status):
-        assert not self.sensor.poke({})
-
-    @mock.patch.object(AthenaHook, "poll_query_status", side_effect=("FAILED",))
-    def test_poke_failed(self, mock_poll_query_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke({})
-        assert "Athena sensor failed" in str(ctx.value)
-
-    @mock.patch.object(AthenaHook, "poll_query_status", side_effect=("CANCELLED",))
-    def test_poke_cancelled(self, mock_poll_query_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke({})
-        assert "Athena sensor failed" in str(ctx.value)
+    @pytest.mark.parametrize("poll_query_status", ["SUCCEEDED"])
+    def test_poke_true_on_status(self, poll_query_status):
+        with mock.patch.object(AthenaHook, "poll_query_status", side_effect=[poll_query_status]):
+            assert self.sensor.poke({})
+
+    @pytest.mark.parametrize("poll_query_status", ["RUNNING", "QUEUED"])
+    def test_poke_false_on_status(self, poll_query_status):
+        with mock.patch.object(AthenaHook, "poll_query_status", side_effect=[poll_query_status]):
+            assert not self.sensor.poke({})
+
+    @pytest.mark.parametrize("poll_query_status", ["FAILED", "CANCELLED"])
+    def test_poke_raise_on_status(self, poll_query_status):
+        with mock.patch.object(AthenaHook, "poll_query_status", side_effect=[poll_query_status]):
+            with pytest.raises(AirflowException, match=r"Athena sensor failed"):
+                self.sensor.poke({})


Merge same tests into parametrize tests

Taragolis · 2022-11-29T04:39:48Z

tests/providers/amazon/aws/sensors/test_dms_task.py

-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("stopped",))
-    def test_poke_stopped(self, mock_get_task_status):
-        assert self.sensor.poke(None)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("running",))
-    def test_poke_running(self, mock_get_task_status):
-        assert not self.sensor.poke(None)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("starting",))
-    def test_poke_starting(self, mock_get_task_status):
-        assert not self.sensor.poke(None)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("ready",))
-    def test_poke_ready(self, mock_get_task_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "Unexpected status: ready" in str(ctx.value)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("creating",))
-    def test_poke_creating(self, mock_get_task_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "Unexpected status: creating" in str(ctx.value)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("failed",))
-    def test_poke_failed(self, mock_get_task_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "Unexpected status: failed" in str(ctx.value)
-
-    @mock.patch.object(DmsHook, "get_task_status", side_effect=("deleting",))
-    def test_poke_deleting(self, mock_get_task_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "Unexpected status: deleting" in str(ctx.value)
+    @pytest.mark.parametrize("task_status", ["stopped"])
+    def test_poke_true_on_status(self, task_status):
+        with mock.patch.object(DmsHook, "get_task_status", side_effect=[task_status]):
+            assert self.sensor.poke({})
+
+    @pytest.mark.parametrize("task_status", ["running", "starting"])
+    def test_poke_false_on_status(self, task_status):
+        with mock.patch.object(DmsHook, "get_task_status", side_effect=[task_status]):
+            assert not self.sensor.poke({})
+
+    @pytest.mark.parametrize("task_status", ["ready", "creating", "failed", "deleting"])
+    def test_poke_raise_unexpected_status_on_status(self, task_status):
+        with mock.patch.object(DmsHook, "get_task_status", side_effect=[task_status]):
+            with pytest.raises(AirflowException, match=rf"Unexpected status: {task_status}"):
+                self.sensor.poke({})


Merge same tests into parametrize tests

Taragolis · 2022-11-29T04:39:58Z

tests/providers/amazon/aws/sensors/test_emr_containers.py

-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("PENDING",))
-    def test_poke_pending(self, mock_check_query_status):
-        assert not self.sensor.poke(None)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("SUBMITTED",))
-    def test_poke_submitted(self, mock_check_query_status):
-        assert not self.sensor.poke(None)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("RUNNING",))
-    def test_poke_running(self, mock_check_query_status):
-        assert not self.sensor.poke(None)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("COMPLETED",))
-    def test_poke_completed(self, mock_check_query_status):
-        assert self.sensor.poke(None)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("FAILED",))
-    def test_poke_failed(self, mock_check_query_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "EMR Containers sensor failed" in str(ctx.value)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("CANCELLED",))
-    def test_poke_cancelled(self, mock_check_query_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "EMR Containers sensor failed" in str(ctx.value)
-
-    @mock.patch.object(EmrContainerHook, "check_query_status", side_effect=("CANCEL_PENDING",))
-    def test_poke_cancel_pending(self, mock_check_query_status):
-        with pytest.raises(AirflowException) as ctx:
-            self.sensor.poke(None)
-        assert "EMR Containers sensor failed" in str(ctx.value)
+    @pytest.mark.parametrize("query_status", ["COMPLETED"])
+    def test_poke_true_on_query_status(self, query_status):
+        with mock.patch.object(EmrContainerHook, "check_query_status", side_effect=[query_status]):
+            assert self.sensor.poke({})
+
+    @pytest.mark.parametrize("query_status", ["PENDING", "SUBMITTED", "RUNNING"])
+    def test_poke_false_on_query_status(self, query_status):
+        with mock.patch.object(EmrContainerHook, "check_query_status", side_effect=[query_status]):
+            assert not self.sensor.poke({})
+
+    @pytest.mark.parametrize("query_status", ["FAILED", "CANCELLED", "CANCEL_PENDING"])
+    def test_poke_raise_on_query_status(self, query_status):
+        with mock.patch.object(EmrContainerHook, "check_query_status", side_effect=[query_status]):
+            with pytest.raises(AirflowException, match=r"EMR Containers sensor failed"):
+                assert not self.sensor.poke({})


Merge same tests into parametrize tests

Taragolis · 2022-11-29T04:42:04Z

tests/providers/amazon/aws/operators/test_sagemaker_training.py

+    def setup_method(self):
+        self.create_training_params = copy.deepcopy(CREATE_TRAINING_PARAMS)
        self.sagemaker = SageMakerTrainingOperator(
            task_id="test_sagemaker_operator",
-            config=CREATE_TRAINING_PARAMS,
+            config=self.create_training_params,
            wait_for_completion=False,
            check_interval=5,
        )


After migrate to pytest some test failed due to mutability of test parameters so this parameters recreate for each test case.

Taragolis · 2022-11-29T04:57:48Z

tests/providers/amazon/aws/hooks/test_emr_containers.py

-        query_status = self.emr_containers.poll_query_status(job_id="job123456", max_polling_attempts=2)
+        query_status = self.emr_containers.poll_query_status(
+            job_id="job123456", max_polling_attempts=2, poll_interval=2
+        )


This simple change actually speedup test execution from 30 sec to 2 seconds

before

============================ slowest 100 durations ============================= 30.03s call tests/providers/amazon/aws/hooks/test_emr_containers.py::TestEmrContainerHook::test_query_status_polling_with_timeout

after

2.01s call tests/providers/amazon/aws/hooks/test_emr_containers.py::TestEmrContainerHook::test_query_status_polling_with_timeout

Taragolis · 2022-12-01T18:24:29Z

I decide split this changes to multiple separate PR's

Could be really reviewed
Less potential merge conflicts

boring-cyborg bot added area:providers provider:amazon AWS/Amazon - related issues labels Nov 28, 2022

Taragolis force-pushed the amazon-provider-tests-convert-pytest branch from e9fbf46 to 191a6e6 Compare November 28, 2022 19:52

Replace unittests in amazon provider tests by pure pytest

48ce15b

Taragolis force-pushed the amazon-provider-tests-convert-pytest branch from 191a6e6 to 48ce15b Compare November 28, 2022 19:54

add missing positional method argument

6ecb2f8

o-nikolas approved these changes Nov 29, 2022

View reviewed changes

Taragolis commented Nov 29, 2022

View reviewed changes

Taragolis marked this pull request as ready for review November 29, 2022 04:45

Taragolis requested a review from eladkal as a code owner November 29, 2022 04:45

Taragolis commented Nov 29, 2022

View reviewed changes

Taragolis added the changelog:skip Changes that should be skipped from the changelog (CI, tests, etc..) label Nov 30, 2022

Taragolis closed this Dec 1, 2022

Taragolis mentioned this pull request Dec 1, 2022

Migrate amazon provider hooks tests from unittests to pytest #28039

Merged

Taragolis deleted the amazon-provider-tests-convert-pytest branch January 14, 2023 18:29

Taragolis restored the amazon-provider-tests-convert-pytest branch January 14, 2023 18:29

		mock_create_fargate_profile.assert_called_with(**convert_keys(parameters))


		class TestEksCreateNodegroupOperator:

Replace unittests in amazon provider tests by pure pytest #27970

Replace unittests in amazon provider tests by pure pytest #27970

Uh oh!

Conversation

Taragolis commented Nov 28, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

o-nikolas left a comment

Choose a reason for hiding this comment

Uh oh!

Taragolis left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Taragolis Nov 29, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Taragolis commented Dec 1, 2022

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Replace `unittests` in amazon provider tests by pure `pytest` #27970

Replace `unittests` in amazon provider tests by pure `pytest` #27970

Taragolis commented Nov 28, 2022 •

edited

Loading

Taragolis Nov 29, 2022 •

edited

Loading