Skip to content

Commit

Permalink
Adding retry logic for #655 and add tests for delete_default_vpc.py (#…
Browse files Browse the repository at this point in the history
…708)

* Adding retry logic for #655 and add tests for delete_default_vpc.py

* Proper logging

* Adopt tenacity in favor of retrying

* Max retry logic and exponential backoff intervals

* Change logging msg to DEBUG and add meaningfull msg

* Mock nap.time.sleep for faster tests

* Update src/lambda_codebase/account_processing/tests/test_delete_default_vpc.py

---------

Co-authored-by: Javy de Koning
Co-authored-by: Stewart Wallace
Co-authored-by: Simon Kok
  • Loading branch information
javydekoning committed Apr 19, 2024
1 parent 325c960 commit 32e9ecf
Show file tree
Hide file tree
Showing 4 changed files with 79 additions and 5 deletions.
26 changes: 21 additions & 5 deletions src/lambda_codebase/account_processing/delete_default_vpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,10 @@
Deletes the default VPC in a particular region
"""
import os

from aws_xray_sdk.core import patch_all
from botocore.exceptions import ClientError
import tenacity

# ADF imports
from logger import configure_logger
Expand All @@ -25,12 +28,25 @@ def assume_role(account_id):
"adf_delete_default_vpc",
)


@tenacity.retry(
retry=tenacity.retry_if_exception_type(ClientError),
# Fail after 180 Sec of retrying
stop=tenacity.stop_after_delay(180),
# Wait 2^x * 1 second between each retry starting with 4s, max 10s intervals
wait=tenacity.wait_exponential(multiplier=1, min=4, max=10),
)
def find_default_vpc(ec2_client):
vpc_response = ec2_client.describe_vpcs()
for vpc in vpc_response["Vpcs"]:
if vpc["IsDefault"] is True:
return vpc["VpcId"]
try:
vpc_response = ec2_client.describe_vpcs()
for vpc in vpc_response["Vpcs"]:
if vpc.get("IsDefault", False):
return vpc["VpcId"]
except ClientError as error:
LOGGER.debug(
"An error occurred while describing VPCs: %s", error
)
raise
# If no default VPC found, return None
return None


Expand Down
1 change: 1 addition & 0 deletions src/lambda_codebase/account_processing/requirements.txt
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
aws-xray-sdk==2.13.0
pyyaml~=6.0.1
tenacity==8.2.3
9 changes: 9 additions & 0 deletions src/lambda_codebase/account_processing/tests/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,3 +2,12 @@
# SPDX-License-Identifier: MIT-0

# pylint: skip-file

"""
__init__ for tests module
"""

import sys
import os

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# Copyright Amazon.com Inc. or its affiliates.
# SPDX-License-Identifier: MIT-0

"""
Tests the delete_default_vpc lambda
"""

import unittest
from unittest.mock import MagicMock, patch
from delete_default_vpc import find_default_vpc
from botocore.exceptions import ClientError


class TestFindDefaultVPC(unittest.TestCase):

@patch("tenacity.nap.time.sleep", MagicMock())
@patch('delete_default_vpc.patch_all')
# pylint: disable=unused-argument
def test_find_default_vpc(self, mock_patch_all):
# Create a mock ec2_client
mock_ec2_client = MagicMock()

# Define the side effects for describe_vpcs method
side_effects = [
ClientError({'Error': {'Code': 'MockTestError'}}, 'describe_vpcs'),
ClientError({'Error': {'Code': 'MockTestError'}}, 'describe_vpcs'),
{"Vpcs": [
{"VpcId": "vpc-123", "IsDefault": False},
{"VpcId": "vpc-456", "IsDefault": True},
{"VpcId": "vpc-789", "IsDefault": False}
]}
]

# Set side_effect for the mock ec2_client.describe_vpcs
mock_ec2_client.describe_vpcs.side_effect = side_effects

# Call the function with the mock ec2_client
default_vpc_id = find_default_vpc(mock_ec2_client)

# Check if the correct default VPC ID is returned
self.assertEqual(default_vpc_id, "vpc-456")

# Check if describe_vpcs method is called 3 times
self.assertEqual(mock_ec2_client.describe_vpcs.call_count, 3)


if __name__ == '__main__':
unittest.main()

0 comments on commit 32e9ecf

Please sign in to comment.