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

src/lambda_codebase/account_processing/delete_default_vpc.py

@@ -5,7 +5,9 @@
 Deletes the default VPC in a particular region
 """
 import os
+import retrying
 from aws_xray_sdk.core import patch_all
+from botocore.exceptions import ClientError
 
 # ADF imports
 from logger import configure_logger
@@ -26,11 +28,23 @@ def assume_role(account_id):
     )
 
 
+@retrying.retry(
+    stop_max_attempt_number=5,
+    wait_exponential_multiplier=10000, # 10 seconds
+    wait_exponential_max=60000 # Maximum backoff time of 60 seconds
+)
 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 e:
+        # Log the error if needed
+        print("An error occurred:", e)
+        # Raise the exception to trigger the retry logic
+        raise
+    # If no default VPC found, return None
     return None
 
 

src/lambda_codebase/account_processing/requirements.txt

@@ -1,2 +1,3 @@
 aws-xray-sdk==2.12.1
 pyyaml~=6.0.1
+retrying~=1.3.4

src/lambda_codebase/account_processing/tests/__init__.py

@@ -0,0 +1,11 @@
+# Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+"""
+__init__ for tests module
+"""
+
+import sys
+import os
+
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

src/lambda_codebase/account_processing/tests/test_delete_default_vpc.py

@@ -0,0 +1,44 @@
+"""
+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('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()