Skip to content

Commit

Permalink
feat(parameters): Allow settings boto3.client() arguments
Browse files Browse the repository at this point in the history
  • Loading branch information
michaelbrewer committed May 1, 2022
1 parent 8ca082f commit d2673d8
Show file tree
Hide file tree
Showing 5 changed files with 147 additions and 12 deletions.
13 changes: 10 additions & 3 deletions aws_lambda_powertools/utilities/parameters/appconfig.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,9 @@ class AppConfigProvider(BaseProvider):
config: botocore.config.Config, optional
Botocore configuration to pass during client initialization
boto3_session : boto3.session.Session, optional
Boto3 session to use for AWS API communication
Boto3 session to use for AWS API communication, will not be used if boto3_client is not None
boto3_client: AppConfigClient, optional
Boto3 Client to use for AWS API communication, will be used instead of boto3_session if both provided
Example
-------
Expand Down Expand Up @@ -68,14 +70,19 @@ def __init__(
application: Optional[str] = None,
config: Optional[Config] = None,
boto3_session: Optional[boto3.session.Session] = None,
boto3_client=None,
):
"""
Initialize the App Config client
"""

config = config or Config()
session = boto3_session or boto3.session.Session()
self.client = session.client("appconfig", config=config)
if boto3_client is not None:
self.client = boto3_client
else:
session = boto3_session or boto3.session.Session()
self.client = session.client("appconfig", config=config)

self.application = resolve_env_var_choice(
choice=application, env=os.getenv(constants.SERVICE_NAME_ENV, "service_undefined")
)
Expand Down
15 changes: 11 additions & 4 deletions aws_lambda_powertools/utilities/parameters/secrets.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ class SecretsProvider(BaseProvider):
config: botocore.config.Config, optional
Botocore configuration to pass during client initialization
boto3_session : boto3.session.Session, optional
Boto3 session to use for AWS API communication
Boto3 session to use for AWS API communication, will not be used if boto3_client is not None
boto3_client: SecretsManagerClient, optional
Boto3 Client to use for AWS API communication, will be used instead of boto3_session if both provided
Example
-------
Expand Down Expand Up @@ -60,14 +62,19 @@ class SecretsProvider(BaseProvider):

client: Any = None

def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
def __init__(
self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None, boto3_client=None
):
"""
Initialize the Secrets Manager client
"""

config = config or Config()
session = boto3_session or boto3.session.Session()
self.client = session.client("secretsmanager", config=config)
if boto3_client is not None:
self.client = boto3_client
else:
session = boto3_session or boto3.session.Session()
self.client = session.client("secretsmanager", config=config)

super().__init__()

Expand Down
15 changes: 11 additions & 4 deletions aws_lambda_powertools/utilities/parameters/ssm.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@ class SSMProvider(BaseProvider):
config: botocore.config.Config, optional
Botocore configuration to pass during client initialization
boto3_session : boto3.session.Session, optional
Boto3 session to use for AWS API communication
Boto3 session to use for AWS API communication, will not be used if boto3_client is not None
boto3_client: SSMClient, optional
Boto3 Client to use for AWS API communication, will be used instead of boto3_session if both provided
Example
-------
Expand Down Expand Up @@ -76,14 +78,19 @@ class SSMProvider(BaseProvider):

client: Any = None

def __init__(self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None):
def __init__(
self, config: Optional[Config] = None, boto3_session: Optional[boto3.session.Session] = None, boto3_client=None
):
"""
Initialize the SSM Parameter Store client
"""

config = config or Config()
session = boto3_session or boto3.session.Session()
self.client = session.client("ssm", config=config)
if boto3_client is not None:
self.client = boto3_client
else:
session = boto3_session or boto3.session.Session()
self.client = session.client("ssm", config=config)

super().__init__()

Expand Down
16 changes: 15 additions & 1 deletion docs/utilities/parameters.md
Original file line number Diff line number Diff line change
Expand Up @@ -482,7 +482,7 @@ Here is the mapping between this utility's functions and methods and the underly

### Customizing boto configuration

The **`config`** and **`boto3_session`** parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) or a custom [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html) when constructing any of the built-in provider classes.
The **`config`** , **`boto3_session`**, and **`boto3_client`** parameters enable you to pass in a custom [botocore config object](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) , [boto3 session](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html), or a [boto3 client](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/boto3.html) when constructing any of the built-in provider classes.

???+ tip
You can use a custom session for retrieving parameters cross-account/region and for snapshot testing.
Expand Down Expand Up @@ -510,6 +510,20 @@ The **`config`** and **`boto3_session`** parameters enable you to pass in a cust
boto_config = Config()
ssm_provider = parameters.SSMProvider(config=boto_config)

def handler(event, context):
# Retrieve a single parameter
value = ssm_provider.get("/my/parameter")
...
```
=== "Custom client"

```python hl_lines="2 4 5"
from aws_lambda_powertools.utilities import parameters
import boto3

boto3_client= session.client(service_name="ssm", endpoint_url='custom_endpoint')
ssm_provider = parameters.SSMProvider(boto3_client=boto3_client)

def handler(event, context):
# Retrieve a single parameter
value = ssm_provider.get("/my/parameter")
Expand Down
100 changes: 100 additions & 0 deletions tests/functional/test_utilities_parameters.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
from io import BytesIO
from typing import Dict

import boto3
import pytest
from boto3.dynamodb.conditions import Key
from botocore import stub
Expand Down Expand Up @@ -444,6 +445,43 @@ def test_ssm_provider_get(mock_name, mock_value, mock_version, config):
stubber.deactivate()


def test_ssm_provider_get_with_custom_client(mock_name, mock_value, mock_version, config):
"""
Test SSMProvider.get() with a non-cached value
"""

client = boto3.client("ssm", config=config)

# Create a new provider
provider = parameters.SSMProvider(boto3_client=client)

# Stub the boto3 client
stubber = stub.Stubber(provider.client)
response = {
"Parameter": {
"Name": mock_name,
"Type": "String",
"Value": mock_value,
"Version": mock_version,
"Selector": f"{mock_name}:{mock_version}",
"SourceResult": "string",
"LastModifiedDate": datetime(2015, 1, 1),
"ARN": f"arn:aws:ssm:us-east-2:111122223333:parameter/{mock_name}",
}
}
expected_params = {"Name": mock_name, "WithDecryption": False}
stubber.add_response("get_parameter", response, expected_params)
stubber.activate()

try:
value = provider.get(mock_name)

assert value == mock_value
stubber.assert_no_pending_responses()
finally:
stubber.deactivate()


def test_ssm_provider_get_default_config(monkeypatch, mock_name, mock_value, mock_version):
"""
Test SSMProvider.get() without specifying the config
Expand Down Expand Up @@ -925,6 +963,37 @@ def test_secrets_provider_get(mock_name, mock_value, config):
stubber.deactivate()


def test_secrets_provider_get_with_custom_client(mock_name, mock_value, config):
"""
Test SecretsProvider.get() with a non-cached value
"""
client = boto3.client("secretsmanager", config=config)

# Create a new provider
provider = parameters.SecretsProvider(boto3_client=client)

# Stub the boto3 client
stubber = stub.Stubber(provider.client)
response = {
"ARN": f"arn:aws:secretsmanager:us-east-1:132456789012:secret/{mock_name}",
"Name": mock_name,
"VersionId": "7a9155b8-2dc9-466e-b4f6-5bc46516c84d",
"SecretString": mock_value,
"CreatedDate": datetime(2015, 1, 1),
}
expected_params = {"SecretId": mock_name}
stubber.add_response("get_secret_value", response, expected_params)
stubber.activate()

try:
value = provider.get(mock_name)

assert value == mock_value
stubber.assert_no_pending_responses()
finally:
stubber.deactivate()


def test_secrets_provider_get_default_config(monkeypatch, mock_name, mock_value):
"""
Test SecretsProvider.get() without specifying a config
Expand Down Expand Up @@ -1555,6 +1624,37 @@ def test_appconf_provider_get_configuration_json_content_type(mock_name, config)
stubber.deactivate()


def test_appconf_provider_get_configuration_json_content_type_with_custom_client(mock_name, config):
"""
Test get_configuration.get with default values
"""

client = boto3.client("appconfig", config=config)

# Create a new provider
environment = "dev"
application = "myapp"
provider = parameters.AppConfigProvider(environment=environment, application=application, boto3_client=client)

mock_body_json = {"myenvvar1": "Black Panther", "myenvvar2": 3}
encoded_message = json.dumps(mock_body_json).encode("utf-8")
mock_value = StreamingBody(BytesIO(encoded_message), len(encoded_message))

# Stub the boto3 client
stubber = stub.Stubber(provider.client)
response = {"Content": mock_value, "ConfigurationVersion": "1", "ContentType": "application/json"}
stubber.add_response("get_configuration", response)
stubber.activate()

try:
value = provider.get(mock_name, transform="json", ClientConfigurationVersion="2")

assert value == mock_body_json
stubber.assert_no_pending_responses()
finally:
stubber.deactivate()


def test_appconf_provider_get_configuration_no_transform(mock_name, config):
"""
Test appconfigprovider.get with default values
Expand Down

0 comments on commit d2673d8

Please sign in to comment.