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(