diff --git a/handler.py b/handler.py index 01b9211..3fa8b14 100644 --- a/handler.py +++ b/handler.py @@ -18,22 +18,28 @@ get_message_for_teams, get_org_message_for_teams, get_message_for_email, get_org_message_for_email, \ get_org_message_for_eventbridge, get_message_for_eventbridge -# query active health API endpoint -health_dns = socket.gethostbyname_ex('global.health.amazonaws.com') -(current_endpoint, global_endpoint, ip_endpoint) = health_dns -health_active_list = current_endpoint.split('.') -health_active_region = health_active_list[1] -print("current health region: ", health_active_region) - -# create a boto3 health client w/ backoff/retry -config = Config( - region_name=health_active_region, - retries=dict( - max_attempts=10 # org view apis have a lower tps than the single - # account apis so we need to use larger - # backoff/retry values than than the boto defaults + +boto3_config = None + + +# Create a boto3 config object +def create_boto3_config(): + # query active health API endpoint + health_dns = socket.gethostbyname_ex('global.health.amazonaws.com') + (current_endpoint, global_endpoint, ip_endpoint) = health_dns + health_active_list = current_endpoint.split('.') + health_active_region = health_active_list[1] + print("current health region: ", health_active_region) + # create a boto3 health client w/ backoff/retry + return Config( + region_name=health_active_region, + retries=dict( + max_attempts=10 # org view apis have a lower tps than the single + # account apis so we need to use larger + # backoff/retry values than than the boto defaults + ) ) -) + # Get Account Name def get_account_name(account_id): @@ -844,20 +850,22 @@ def get_sts_token(service): # create service client using the assumed role credentials, e.g. S3 boto3_client = boto3.client( service, - config=config, + config=boto3_config, aws_access_key_id=ACCESS_KEY, aws_secret_access_key=SECRET_KEY, aws_session_token=SESSION_TOKEN, ) print("Running in member account deployment mode") else: - boto3_client = boto3.client(service, config=config) + boto3_client = boto3.client(service, config=boto3_config) print("Running in management account deployment mode") return boto3_client def main(event, context): print("THANK YOU FOR CHOOSING AWS HEALTH AWARE!") + global boto3_config + boto3_config = create_boto3_config() health_client = get_sts_token('health') org_status = os.environ['ORG_STATUS'] #str_ddb_format_sec = '%s' diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 0000000..9d3ee85 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,3 @@ +boto3==1.18.57 +pytest==6.2.5 +pytest-mock==3.6.1 diff --git a/test_create_boto3_config.py b/test_create_boto3_config.py new file mode 100644 index 0000000..90dcc57 --- /dev/null +++ b/test_create_boto3_config.py @@ -0,0 +1,10 @@ +import handler +import pytest + + +@pytest.mark.parametrize('region', ['us-east-1', 'us-east-2']) +def test_boto3_config_object(region, mocker): + """When the DNS record changes for the health API the config should be updated correctly""" + mock_gethostname = mocker.patch('handler.socket.gethostbyname_ex') + mock_gethostname.return_value = (f'health.{region}.amazonaws.com', ['global.health.amazonaws.com'], ['52.94.233.29']) + assert handler.create_boto3_config().region_name == region