Migrate amazon provider transfer tests from unittests to pytest

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

@@ -42,7 +42,7 @@ def setup_method(self):
     @mock.patch.object(BatchClientHook, "get_job_description")
     def test_poke_on_success_state(self, mock_get_job_description):
         mock_get_job_description.return_value = {"status": "SUCCEEDED"}
-        assert self.batch_sensor.poke({})
+        assert self.batch_sensor.poke({}) is True
         mock_get_job_description.assert_called_once_with(JOB_ID)
 
     @mock.patch.object(BatchClientHook, "get_job_description")

tests/providers/amazon/aws/transfers/test_dynamodb_to_s3.py

@@ -18,14 +18,13 @@
 from __future__ import annotations
 
 import json
-import unittest
 from decimal import Decimal
 from unittest.mock import MagicMock, patch
 
 from airflow.providers.amazon.aws.transfers.dynamodb_to_s3 import DynamoDBToS3Operator, JSONEncoder
 
 
-class JSONEncoderTest(unittest.TestCase):
+class JSONEncoderTest:
     def test_jsonencoder_with_decimal(self):
         """Test JSONEncoder correctly encodes and decodes decimal values."""
 
@@ -36,8 +35,8 @@ def test_jsonencoder_with_decimal(self):
             self.assertAlmostEqual(decoded, org)
 
 
-class DynamodbToS3Test(unittest.TestCase):
-    def setUp(self):
+class DynamodbToS3Test:
+    def setup_method(self):
         self.output_queue = []
 
     def mock_upload_file(self, Filename, Bucket, Key):

tests/providers/amazon/aws/transfers/test_ftp_to_s3.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 from airflow.providers.amazon.aws.transfers.ftp_to_s3 import FTPToS3Operator
@@ -32,7 +31,7 @@
 FTP_PATH_MULTIPLE = "/tmp/"
 
 
-class TestFTPToS3Operator(unittest.TestCase):
+class TestFTPToS3Operator:
     def assert_execute(
         self, mock_local_tmp_file, mock_s3_hook_load_file, mock_ftp_hook_retrieve_file, ftp_file, s3_file
     ):

tests/providers/amazon/aws/transfers/test_google_api_to_s3.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest.mock import Mock, patch
 
 import pytest
@@ -29,8 +28,8 @@
 from airflow.utils import db
 
 
-class TestGoogleApiToS3(unittest.TestCase):
-    def setUp(self):
+class TestGoogleApiToS3:
+    def setup_method(self):
         conf.load_test_config()
 
         db.merge_conn(

tests/providers/amazon/aws/transfers/test_hive_to_dynamodb.py

@@ -19,7 +19,6 @@
 
 import datetime
 import json
-import unittest
 from unittest import mock
 
 import pandas as pd
@@ -34,8 +33,8 @@
 DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10]
 
 
-class TestHiveToDynamoDBOperator(unittest.TestCase):
-    def setUp(self):
+class TestHiveToDynamoDBOperator:
+    def setup_method(self):
         args = {"owner": "airflow", "start_date": DEFAULT_DATE}
         dag = DAG("test_dag_id", default_args=args)
         self.dag = dag

tests/providers/amazon/aws/transfers/test_imap_attachment_to_s3.py

@@ -17,14 +17,13 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
-from unittest.mock import patch
+from unittest import mock
 
 from airflow.providers.amazon.aws.transfers.imap_attachment_to_s3 import ImapAttachmentToS3Operator
 
 
-class TestImapAttachmentToS3Operator(unittest.TestCase):
-    def setUp(self):
+class TestImapAttachmentToS3Operator:
+    def setup_method(self):
         self.kwargs = dict(
             imap_attachment_name="test_file",
             s3_bucket="test_bucket",
@@ -37,8 +36,8 @@ def setUp(self):
             dag=None,
         )
 
-    @patch("airflow.providers.amazon.aws.transfers.imap_attachment_to_s3.S3Hook")
-    @patch("airflow.providers.amazon.aws.transfers.imap_attachment_to_s3.ImapHook")
+    @mock.patch("airflow.providers.amazon.aws.transfers.imap_attachment_to_s3.S3Hook")
+    @mock.patch("airflow.providers.amazon.aws.transfers.imap_attachment_to_s3.ImapHook")
     def test_execute(self, mock_imap_hook, mock_s3_hook):
         mock_imap_hook.return_value.__enter__ = mock_imap_hook
         mock_imap_hook.return_value.retrieve_mail_attachments.return_value = [("test_file", b"Hello World")]

tests/providers/amazon/aws/transfers/test_local_to_s3.py

@@ -19,20 +19,19 @@
 
 import datetime
 import os
-import unittest
 
 import boto3
+import pytest
 from moto import mock_s3
 
 from airflow.models.dag import DAG
 from airflow.providers.amazon.aws.transfers.local_to_s3 import LocalFilesystemToS3Operator
 
+CONFIG = {"verify": False, "replace": False, "encrypt": False, "gzip": False}
 
-class TestFileToS3Operator(unittest.TestCase):
 
-    _config = {"verify": False, "replace": False, "encrypt": False, "gzip": False}
-
-    def setUp(self):
+class TestFileToS3Operator:
+    def setup_method(self):
         args = {"owner": "airflow", "start_date": datetime.datetime(2017, 1, 1)}
         self.dag = DAG("test_dag_id", default_args=args)
         self.dest_key = "test/test1.csv"
@@ -41,7 +40,7 @@ def setUp(self):
         with open(self.testfile1, "wb") as f:
             f.write(b"x" * 393216)
 
-    def tearDown(self):
+    def teardown_method(self):
         os.remove(self.testfile1)
 
     def test_init(self):
@@ -51,15 +50,15 @@ def test_init(self):
             filename=self.testfile1,
             dest_key=self.dest_key,
             dest_bucket=self.dest_bucket,
-            **self._config,
+            **CONFIG,
         )
         assert operator.filename == self.testfile1
         assert operator.dest_key == self.dest_key
         assert operator.dest_bucket == self.dest_bucket
-        assert operator.verify == self._config["verify"]
-        assert operator.replace == self._config["replace"]
-        assert operator.encrypt == self._config["encrypt"]
-        assert operator.gzip == self._config["gzip"]
+        assert operator.verify == CONFIG["verify"]
+        assert operator.replace == CONFIG["replace"]
+        assert operator.encrypt == CONFIG["encrypt"]
+        assert operator.gzip == CONFIG["gzip"]
 
     def test_execute_exception(self):
         operator = LocalFilesystemToS3Operator(
@@ -68,9 +67,9 @@ def test_execute_exception(self):
             filename=self.testfile1,
             dest_key=f"s3://dummy/{self.dest_key}",
             dest_bucket=self.dest_bucket,
-            **self._config,
+            **CONFIG,
         )
-        with self.assertRaises(TypeError):
+        with pytest.raises(TypeError):
             operator.execute(None)
 
     @mock_s3
@@ -83,7 +82,7 @@ def test_execute(self):
             filename=self.testfile1,
             dest_key=self.dest_key,
             dest_bucket=self.dest_bucket,
-            **self._config,
+            **CONFIG,
         )
         operator.execute(None)
 
@@ -102,7 +101,7 @@ def test_execute_with_only_key(self):
             dag=self.dag,
             filename=self.testfile1,
             dest_key=f"s3://dummy/{self.dest_key}",
-            **self._config,
+            **CONFIG,
         )
         operator.execute(None)
 

tests/providers/amazon/aws/transfers/test_mongo_to_s3.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 from airflow.models import DAG, DagRun, TaskInstance
@@ -40,8 +39,8 @@
 ]
 
 
-class TestMongoToS3Operator(unittest.TestCase):
-    def setUp(self):
+class TestMongoToS3Operator:
+    def setup_method(self):
         args = {"owner": "airflow", "start_date": DEFAULT_DATE}
 
         self.dag = DAG("test_dag_id", default_args=args)

tests/providers/amazon/aws/transfers/test_redshift_to_s3.py

@@ -17,37 +17,31 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
+import pytest
 from boto3.session import Session
-from parameterized import parameterized
 
 from airflow.models.connection import Connection
 from airflow.providers.amazon.aws.transfers.redshift_to_s3 import RedshiftToS3Operator
 from airflow.providers.amazon.aws.utils.redshift import build_credentials_block
 from tests.test_utils.asserts import assert_equal_ignore_multiple_spaces
 
 
-class TestRedshiftToS3Transfer(unittest.TestCase):
-    @parameterized.expand(
-        [
-            [True, "key/table_"],
-            [False, "key"],
-        ]
-    )
+class TestRedshiftToS3Transfer:
+    @pytest.mark.parametrize("table_as_file_name, expected_s3_key", [[True, "key/table_"], [False, "key"]])
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
     @mock.patch("airflow.models.connection.Connection")
     @mock.patch("boto3.session.Session")
     @mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook.run")
     def test_table_unloading(
         self,
-        table_as_file_name,
-        expected_s3_key,
         mock_run,
         mock_session,
         mock_connection,
         mock_hook,
+        table_as_file_name,
+        expected_s3_key,
     ):
         access_key = "aws_access_key_id"
         secret_key = "aws_secret_access_key"
@@ -94,24 +88,19 @@ def test_table_unloading(
         assert secret_key in unload_query
         assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
 
-    @parameterized.expand(
-        [
-            [True, "key/table_"],
-            [False, "key"],
-        ]
-    )
+    @pytest.mark.parametrize("table_as_file_name, expected_s3_key", [[True, "key/table_"], [False, "key"]])
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
     @mock.patch("airflow.models.connection.Connection")
     @mock.patch("boto3.session.Session")
     @mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook.run")
     def test_execute_sts_token(
         self,
-        table_as_file_name,
-        expected_s3_key,
         mock_run,
         mock_session,
         mock_connection,
         mock_hook,
+        table_as_file_name,
+        expected_s3_key,
     ):
         access_key = "ASIA_aws_access_key_id"
         secret_key = "aws_secret_access_key"
@@ -160,27 +149,28 @@ def test_execute_sts_token(
         assert token in unload_query
         assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
 
-    @parameterized.expand(
+    @pytest.mark.parametrize(
+        "table, table_as_file_name, expected_s3_key",
         [
             ["table", True, "key/table_"],
             ["table", False, "key"],
             [None, False, "key"],
             [None, True, "key"],
-        ]
+        ],
     )
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
     @mock.patch("airflow.models.connection.Connection")
     @mock.patch("boto3.session.Session")
     @mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook.run")
     def test_custom_select_query_unloading(
         self,
-        table,
-        table_as_file_name,
-        expected_s3_key,
         mock_run,
         mock_session,
         mock_connection,
         mock_hook,
+        table,
+        table_as_file_name,
+        expected_s3_key,
     ):
         access_key = "aws_access_key_id"
         secret_key = "aws_secret_access_key"
@@ -225,24 +215,19 @@ def test_custom_select_query_unloading(
         assert secret_key in unload_query
         assert_equal_ignore_multiple_spaces(self, mock_run.call_args[0][0], unload_query)
 
-    @parameterized.expand(
-        [
-            [True, "key/table_"],
-            [False, "key"],
-        ]
-    )
+    @pytest.mark.parametrize("table_as_file_name, expected_s3_key", [[True, "key/table_"], [False, "key"]])
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
     @mock.patch("airflow.models.connection.Connection")
     @mock.patch("boto3.session.Session")
     @mock.patch("airflow.providers.amazon.aws.hooks.redshift_sql.RedshiftSQLHook.run")
     def test_table_unloading_role_arn(
         self,
-        table_as_file_name,
-        expected_s3_key,
         mock_run,
         mock_session,
         mock_connection,
         mock_hook,
+        table_as_file_name,
+        expected_s3_key,
     ):
         access_key = "aws_access_key_id"
         secret_key = "aws_secret_access_key"

tests/providers/amazon/aws/transfers/test_s3_to_ftp.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 from airflow.providers.amazon.aws.transfers.s3_to_ftp import S3ToFTPOperator
@@ -30,7 +29,7 @@
 FTP_CONN_ID = "ftp_default"
 
 
-class TestS3ToFTPOperator(unittest.TestCase):
+class TestS3ToFTPOperator:
     @mock.patch("airflow.providers.ftp.hooks.ftp.FTPHook.store_file")
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_key")
     @mock.patch("airflow.providers.amazon.aws.transfers.s3_to_ftp.NamedTemporaryFile")

tests/providers/amazon/aws/transfers/test_s3_to_redshift.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import unittest
 from unittest import mock
 
 import pytest
@@ -29,7 +28,7 @@
 from tests.test_utils.asserts import assert_equal_ignore_multiple_spaces
 
 
-class TestS3ToRedshiftTransfer(unittest.TestCase):
+class TestS3ToRedshiftTransfer:
     @mock.patch("airflow.providers.amazon.aws.hooks.s3.S3Hook.get_connection")
     @mock.patch("airflow.models.connection.Connection")
     @mock.patch("boto3.session.Session")