From 082b142c55ab4723793a2f21a1487f8c249408b4 Mon Sep 17 00:00:00 2001 From: pnadolny13 Date: Tue, 11 Apr 2023 15:47:33 -0400 Subject: [PATCH] refactor to boto authenticator class --- .../{aws_auth.py => aws_authenticators.py} | 27 ++++++++-- tap_dynamodb/aws_base.py | 32 ------------ tap_dynamodb/dynamo.py | 4 +- ...aws_auth.py => test_aws_authenticators.py} | 52 +++++++++++-------- 4 files changed, 56 insertions(+), 59 deletions(-) rename tap_dynamodb/{aws_auth.py => aws_authenticators.py} (82%) delete mode 100644 tap_dynamodb/aws_base.py rename tests/{test_aws_auth.py => test_aws_authenticators.py} (73%) diff --git a/tap_dynamodb/aws_auth.py b/tap_dynamodb/aws_authenticators.py similarity index 82% rename from tap_dynamodb/aws_auth.py rename to tap_dynamodb/aws_authenticators.py index 8b71138..74b33af 100644 --- a/tap_dynamodb/aws_auth.py +++ b/tap_dynamodb/aws_authenticators.py @@ -4,8 +4,13 @@ import boto3 -class AWSAuth: - def __init__(self, config): +class AWSBotoAuthenticator: + def __init__(self, config, service_name): + self._service_name = service_name + self._config = config + self._client = None + self._resource = None + self.logger = logging.getLogger(__name__) # config for use environment variables if config.get("use_aws_env_vars"): self.aws_access_key_id = os.environ.get("AWS_ACCESS_KEY_ID") @@ -23,7 +28,23 @@ def __init__(self, config): self.aws_endpoint_url = config.get("aws_endpoint_url") self.aws_assume_role_arn = config.get("aws_assume_role_arn") - self.logger = logging.getLogger(__name__) + @property + def client(self): + if self._client: + return self._client + else: + session = self.get_session() + self._client = self.get_client(session, self._service_name) + return self._client + + @property + def resource(self): + if self._resource: + return self._resource + else: + session = self.get_session() + self._resource = self.get_resource(session, self._service_name) + return self._resource def get_session(self): session = None diff --git a/tap_dynamodb/aws_base.py b/tap_dynamodb/aws_base.py deleted file mode 100644 index d15765d..0000000 --- a/tap_dynamodb/aws_base.py +++ /dev/null @@ -1,32 +0,0 @@ -import logging - -from tap_dynamodb.aws_auth import AWSAuth - - -class AWSBase: - def __init__(self, config, service_name): - self._service_name = service_name - self._config = config - self._client = None - self._resource = None - self.logger = logging.getLogger(__name__) - - @property - def client(self): - if self._client: - return self._client - else: - auth_obj = AWSAuth(self._config) - session = auth_obj.get_session() - self._client = auth_obj.get_client(session, self._service_name) - return self._client - - @property - def resource(self): - if self._resource: - return self._resource - else: - auth_obj = AWSAuth(self._config) - session = auth_obj.get_session() - self._resource = auth_obj.get_resource(session, self._service_name) - return self._resource diff --git a/tap_dynamodb/dynamo.py b/tap_dynamodb/dynamo.py index 04fec9c..1e86112 100644 --- a/tap_dynamodb/dynamo.py +++ b/tap_dynamodb/dynamo.py @@ -2,11 +2,11 @@ import orjson from botocore.exceptions import ClientError -from tap_dynamodb.aws_base import AWSBase +from tap_dynamodb.aws_authenticators import AWSBotoAuthenticator from tap_dynamodb.exception import EmptyTableException -class DynamoDB(AWSBase): +class DynamoDB(AWSBotoAuthenticator): def __init__(self, config): super().__init__(config, "dynamodb") diff --git a/tests/test_aws_auth.py b/tests/test_aws_authenticators.py similarity index 73% rename from tests/test_aws_auth.py rename to tests/test_aws_authenticators.py index 18403b4..adc64b7 100644 --- a/tests/test_aws_auth.py +++ b/tests/test_aws_authenticators.py @@ -1,21 +1,22 @@ import pytest -from contextlib import nullcontext as does_not_raise -from unittest import mock from moto import mock_dynamodb, mock_sts -from tap_dynamodb.aws_auth import AWSAuth -import os +from tap_dynamodb.aws_authenticators import AWSBotoAuthenticator from unittest.mock import patch -import botocore -@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session") + +@patch( + "tap_dynamodb.aws_authenticators.boto3.Session", + return_value="mock_session", +) @mock_dynamodb def test_get_session_base(patch): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_access_key_id": "foo", "aws_secret_access_key": "bar", "aws_default_region": "baz", - } + }, + "dynamodb", ) session = auth.get_session() patch.assert_called_with( @@ -25,16 +26,18 @@ def test_get_session_base(patch): ) assert session == "mock_session" -@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session") + +@patch("tap_dynamodb.aws_authenticators.boto3.Session", return_value="mock_session") @mock_dynamodb def test_get_session_w_token(patch): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_access_key_id": "foo", "aws_secret_access_key": "bar", "aws_session_token": "abc", "aws_default_region": "baz", - } + }, + "dynamodb", ) session = auth.get_session() patch.assert_called_with( @@ -46,13 +49,14 @@ def test_get_session_w_token(patch): assert session == "mock_session" -@patch("tap_dynamodb.aws_auth.boto3.Session", return_value="mock_session") +@patch("tap_dynamodb.aws_authenticators.boto3.Session", return_value="mock_session") @mock_dynamodb def test_get_session_w_profile(patch): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_profile": "foo", - } + }, + "dynamodb", ) session = auth.get_session() patch.assert_called_with( @@ -64,44 +68,48 @@ def test_get_session_w_profile(patch): @mock_dynamodb def test_get_session_empty_fail(): with pytest.raises(Exception): - auth = AWSAuth({}) + auth = AWSBotoAuthenticator({}) auth.get_session() @mock_dynamodb @mock_sts def test_get_session_assume_role(): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_access_key_id": "foo", "aws_secret_access_key": "bar", "aws_default_region": "baz", - "aws_assume_role_arn": "arn:aws:iam::123456778910:role/my-role-name" - } + "aws_assume_role_arn": "arn:aws:iam::123456778910:role/my-role-name", + }, + "dynamodb", ) session = auth.get_session() @mock_dynamodb def test_get_client(): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_access_key_id": "foo", "aws_secret_access_key": "bar", "aws_default_region": "baz", - } + }, + "dynamodb", ) session = auth.get_session() client = auth.get_client(session, "dynamodb") + @mock_dynamodb def test_get_resource(): - auth = AWSAuth( + auth = AWSBotoAuthenticator( { "aws_access_key_id": "foo", "aws_secret_access_key": "bar", "aws_default_region": "baz", - } + }, + "dynamodb", ) session = auth.get_session() resource = auth.get_resource(session, "dynamodb")