You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by pi...@apache.org on 2022/10/21 17:20:24 UTC

[airflow] branch main updated: Make amazon provider tests compatible with latest `moto` (#27177)

This is an automated email from the ASF dual-hosted git repository.

pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 804aeee9d7 Make amazon provider tests compatible with latest `moto` (#27177)
804aeee9d7 is described below

commit 804aeee9d791cc39dbd10c5d3248bb9200b119a6
Author: Andrey Anshin <An...@taragol.is>
AuthorDate: Fri Oct 21 21:20:15 2022 +0400

    Make amazon provider tests compatible with latest `moto` (#27177)
---
 tests/providers/amazon/aws/hooks/test_base_aws.py  | 30 +++++++++-------------
 tests/providers/amazon/aws/hooks/test_dynamodb.py  | 13 +++-------
 tests/providers/amazon/aws/hooks/test_eks.py       | 13 +++++++---
 .../amazon/aws/transfers/test_hive_to_dynamodb.py  | 15 +++--------
 4 files changed, 29 insertions(+), 42 deletions(-)

diff --git a/tests/providers/amazon/aws/hooks/test_base_aws.py b/tests/providers/amazon/aws/hooks/test_base_aws.py
index 07e1cc7edc..d7b48b5ece 100644
--- a/tests/providers/amazon/aws/hooks/test_base_aws.py
+++ b/tests/providers/amazon/aws/hooks/test_base_aws.py
@@ -28,7 +28,13 @@ import boto3
 import pytest
 from botocore.config import Config
 from botocore.credentials import ReadOnlyCredentials
-from moto.core import ACCOUNT_ID
+
+try:
+    from moto.core import DEFAULT_ACCOUNT_ID
+except ImportError:
+    from moto.core import ACCOUNT_ID as DEFAULT_ACCOUNT_ID
+
+from moto import mock_dynamodb, mock_emr, mock_iam, mock_sts
 
 from airflow.models.connection import Connection
 from airflow.providers.amazon.aws.hooks.base_aws import (
@@ -39,14 +45,6 @@ from airflow.providers.amazon.aws.hooks.base_aws import (
 from airflow.providers.amazon.aws.utils.connection_wrapper import AwsConnectionWrapper
 from tests.test_utils.config import conf_vars
 
-try:
-    from moto import mock_dynamodb2, mock_emr, mock_iam, mock_sts
-except ImportError:
-    mock_emr = None
-    mock_dynamodb2 = None
-    mock_sts = None
-    mock_iam = None
-
 MOCK_AWS_CONN_ID = "mock-conn-id"
 MOCK_CONN_TYPE = "aws"
 MOCK_BOTO3_SESSION = mock.MagicMock(return_value="Mock boto3.session.Session")
@@ -110,7 +108,7 @@ SAML_ASSERTION = """
     </AuthnStatement>
   </Assertion>
 </samlp:Response>""".format(  # noqa: E501
-    account_id=ACCOUNT_ID,
+    account_id=DEFAULT_ACCOUNT_ID,
     role_name="test-role",
     provider_name="TestProvFed",
     username="testuser",
@@ -241,8 +239,7 @@ class TestAwsBaseHook:
 
         assert client_from_hook.list_clusters()['Clusters'] == []
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_get_resource_type_set_in_class_attribute(self):
         hook = AwsBaseHook(aws_conn_id='aws_default', resource_type='dynamodb')
         resource_from_hook = hook.get_resource_type()
@@ -261,8 +258,7 @@ class TestAwsBaseHook:
 
         assert table.item_count == 0
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_get_session_returns_a_boto3_session(self):
         hook = AwsBaseHook(aws_conn_id='aws_default', resource_type='dynamodb')
         session_from_hook = hook.get_session()
@@ -593,8 +589,7 @@ class TestAwsBaseHook:
             assert mock_refresh.call_count == 2
             assert len(expire_on_calls) == 0
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     @pytest.mark.parametrize("conn_type", ["client", "resource"])
     @pytest.mark.parametrize(
         "connection_uri,region_name,env_region,expected_region_name",
@@ -622,8 +617,7 @@ class TestAwsBaseHook:
 
             assert hook.conn_region_name == expected_region_name
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamo2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     @pytest.mark.parametrize("conn_type", ["client", "resource"])
     @pytest.mark.parametrize(
         "connection_uri,expected_partition",
diff --git a/tests/providers/amazon/aws/hooks/test_dynamodb.py b/tests/providers/amazon/aws/hooks/test_dynamodb.py
index cf5ba45597..a3542b5c31 100644
--- a/tests/providers/amazon/aws/hooks/test_dynamodb.py
+++ b/tests/providers/amazon/aws/hooks/test_dynamodb.py
@@ -20,23 +20,18 @@ from __future__ import annotations
 import unittest
 import uuid
 
-from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
+from moto import mock_dynamodb
 
-try:
-    from moto import mock_dynamodb2
-except ImportError:
-    mock_dynamodb2 = None
+from airflow.providers.amazon.aws.hooks.dynamodb import DynamoDBHook
 
 
 class TestDynamoDBHook(unittest.TestCase):
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_get_conn_returns_a_boto3_connection(self):
         hook = DynamoDBHook(aws_conn_id='aws_default')
         assert hook.get_conn() is not None
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_insert_batch_items_dynamodb_table(self):
 
         hook = DynamoDBHook(
diff --git a/tests/providers/amazon/aws/hooks/test_eks.py b/tests/providers/amazon/aws/hooks/test_eks.py
index 3ecb265580..623cd41527 100644
--- a/tests/providers/amazon/aws/hooks/test_eks.py
+++ b/tests/providers/amazon/aws/hooks/test_eks.py
@@ -29,7 +29,12 @@ import yaml
 from _pytest._code import ExceptionInfo
 from botocore.exceptions import ClientError
 from freezegun import freeze_time
-from moto.core import ACCOUNT_ID
+
+try:
+    from moto.core import DEFAULT_ACCOUNT_ID
+except ImportError:
+    from moto.core import ACCOUNT_ID as DEFAULT_ACCOUNT_ID
+
 from moto.core.exceptions import AWSError
 from moto.eks.exceptions import (
     InvalidParameterException,
@@ -294,7 +299,7 @@ class TestEksHooks:
         expected_arn_values: list = [
             PARTITION,
             REGION,
-            ACCOUNT_ID,
+            DEFAULT_ACCOUNT_ID,
             generated_test_data.cluster_names,
         ]
 
@@ -509,7 +514,7 @@ class TestEksHooks:
         expected_arn_values: list = [
             PARTITION,
             REGION,
-            ACCOUNT_ID,
+            DEFAULT_ACCOUNT_ID,
             generated_test_data.cluster_name,
             generated_test_data.nodegroup_names,
             None,
@@ -911,7 +916,7 @@ class TestEksHooks:
         expected_arn_values: list = [
             PARTITION,
             REGION,
-            ACCOUNT_ID,
+            DEFAULT_ACCOUNT_ID,
             generated_test_data.cluster_name,
             generated_test_data.fargate_profile_names,
             None,
diff --git a/tests/providers/amazon/aws/transfers/test_hive_to_dynamodb.py b/tests/providers/amazon/aws/transfers/test_hive_to_dynamodb.py
index 511bb7aff3..4535141318 100644
--- a/tests/providers/amazon/aws/transfers/test_hive_to_dynamodb.py
+++ b/tests/providers/amazon/aws/transfers/test_hive_to_dynamodb.py
@@ -23,6 +23,7 @@ import unittest
 from unittest import mock
 
 import pandas as pd
+from moto import mock_dynamodb
 
 import airflow.providers.amazon.aws.transfers.hive_to_dynamodb
 from airflow.models.dag import DAG
@@ -32,11 +33,6 @@ DEFAULT_DATE = datetime.datetime(2015, 1, 1)
 DEFAULT_DATE_ISO = DEFAULT_DATE.isoformat()
 DEFAULT_DATE_DS = DEFAULT_DATE_ISO[:10]
 
-try:
-    from moto import mock_dynamodb2
-except ImportError:
-    mock_dynamodb2 = None
-
 
 class TestHiveToDynamoDBOperator(unittest.TestCase):
     def setUp(self):
@@ -50,8 +46,7 @@ class TestHiveToDynamoDBOperator(unittest.TestCase):
     def process_data(data, *args, **kwargs):
         return json.loads(data.to_json(orient='records'))
 
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_get_conn_returns_a_boto3_connection(self):
         hook = DynamoDBHook(aws_conn_id='aws_default')
         assert hook.get_conn() is not None
@@ -60,8 +55,7 @@ class TestHiveToDynamoDBOperator(unittest.TestCase):
         'airflow.providers.apache.hive.hooks.hive.HiveServer2Hook.get_pandas_df',
         return_value=pd.DataFrame(data=[('1', 'sid')], columns=['id', 'name']),
     )
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_get_records_with_schema(self, mock_get_pandas_df):
         # this table needs to be created in production
         self.hook.get_conn().create_table(
@@ -91,8 +85,7 @@ class TestHiveToDynamoDBOperator(unittest.TestCase):
         'airflow.providers.apache.hive.hooks.hive.HiveServer2Hook.get_pandas_df',
         return_value=pd.DataFrame(data=[('1', 'sid'), ('1', 'gupta')], columns=['id', 'name']),
     )
-    @unittest.skipIf(mock_dynamodb2 is None, 'mock_dynamodb2 package not present')
-    @mock_dynamodb2
+    @mock_dynamodb
     def test_pre_process_records_with_schema(self, mock_get_pandas_df):
         # this table needs to be created in production
         self.hook.get_conn().create_table(