Fix zombie snapshots (#802)

cloud_governance/common/clouds/aws/ec2/ec2_operations.py

@@ -341,12 +341,12 @@ def get_volumes(self, **kwargs):
         return self.utils.get_details_resource_list(func_name=self.ec2_client.describe_volumes, input_tag='Volumes',
                                                     check_tag='NextToken', **kwargs)
 
-    def get_images(self):
+    def get_images(self, **kwargs):
         """
         This method returns all images in the region
         @return:
         """
-        return self.ec2_client.describe_images(Owners=['self'])['Images']
+        return self.ec2_client.describe_images(Owners=['self'], **kwargs)['Images']
 
     def get_snapshots(self):
         """

cloud_governance/policy/aws/zombie_snapshots.py

@@ -12,33 +12,20 @@ def __init__(self):
         super().__init__()
         self.__image_ids = self._get_ami_ids()
 
-    def __get_image_ids_from_description(self, snapshot_description: str):
+    def __snapshot_id_in_images(self, resource_id: str):
         """
-        This method gets image Ids from snapshot description
-        Two cases:
-        # Created by CreateImage(i-******) for ami-********
-        # Copied for DestinationAmi ami-******* from SourceAmi ami-******* for SourceSnapshot snap-******. Task created on 1,566,308,778,174.
-        @return:
-        """
-        image_ids = []
-        images_array = snapshot_description.split('ami-')[1:]
-        for image in images_array:
-            image_ids.append(f'ami-{image.split(" ")[0]}')
-        return image_ids
-
-    def __is_zombie_snapshot(self, snapshot_description: str):
-        """
-        This method returns bool on verifying snapshots as zombie or not
-        :param snapshot_description:
+        This method checks if snapshot_id exists in images list
+        :param resource_id:
         :return:
         """
-        zombie_snapshot = True
-        if snapshot_description:
-            snapshot_images = self.__get_image_ids_from_description(snapshot_description)
-            for snapshot_image in snapshot_images:
-                if snapshot_image in self.__image_ids:
-                    return False
-        return zombie_snapshot
+        image_snapshot_filter = {
+            'Name': 'block-device-mapping.snapshot-id',
+            'Values': [
+                resource_id,
+            ]
+        }
+        images = self._get_ami_ids(Filters=[image_snapshot_filter])
+        return len(images) >= 1
 
     def run(self):
         """
@@ -53,7 +40,7 @@ def run(self):
             cleanup_result = False
             cluster_tag = self._get_cluster_tag(tags=tags)
             cleanup_days = 0
-            if not cluster_tag and self.__is_zombie_snapshot(snapshot.get('Description')):
+            if not cluster_tag and not self.__snapshot_id_in_images(resource_id):
                 cleanup_days = self.get_clean_up_days_count(tags=tags)
                 cleanup_result = self.verify_and_delete_resource(resource_id=resource_id, tags=tags,
                                                                  clean_up_days=cleanup_days)

cloud_governance/policy/helpers/aws/aws_policy_operations.py

@@ -289,12 +289,12 @@ def get_network_out_kib_metric(self, resource_id: str, days: int = INSTANCE_IDLE
                                                                          aggregation='average')
         return round(average_network_out_bytes / TOTAL_BYTES_IN_KIB, DEFAULT_ROUND_DIGITS)
 
-    def _get_ami_ids(self):
+    def _get_ami_ids(self, **kwargs):
         """
         This method returns all image ids
         @return:
         """
-        images = self._ec2_operations.get_images()
+        images = self._ec2_operations.get_images(**kwargs)
         image_ids = []
         for image in images:
             image_ids.append(image.get('ImageId'))

tests/unittest/cloud_governance/policy/aws/test_zombie_snapshots.py

@@ -1,5 +1,5 @@
-import os
 from datetime import datetime
+from unittest.mock import patch
 
 import boto3
 from moto import mock_ec2
@@ -10,8 +10,6 @@
 from tests.unittest.configs import DRY_RUN_YES, AWS_DEFAULT_REGION, INSTANCE_TYPE_T2_MICRO, DEFAULT_AMI_ID, \
     TEST_USER_NAME, DRY_RUN_NO
 
-os.environ['AWS_DEFAULT_REGION'] = 'us-east-2'
-os.environ['dry_run'] = 'no'
 
 
 @mock_ec2
@@ -181,7 +179,6 @@ def test_zombie_snapshots_contains_cluster_tag():
     assert len(response) == 0
 
 
-@mock_ec2
 def test_zombie_snapshots_no_zombies():
     """
     This method tests snapshot having the active AMI
@@ -191,31 +188,34 @@ def test_zombie_snapshots_no_zombies():
     environment_variables.environment_variables_dict['AWS_DEFAULT_REGION'] = AWS_DEFAULT_REGION
     environment_variables.environment_variables_dict['policy'] = 'zombie_snapshots'
     tags = [{'Key': 'User', 'Value': TEST_USER_NAME}, {'Key': 'policy', 'Value': 'not-delete'},
-            {'Key': 'DaysCount', 'Value': f'{datetime.utcnow().date()}@7'},
-            {'Key': 'kubernetes.io/cluster/test-zombie-cluster', 'Value': f'owned'}]
-    ec2_client = boto3.client('ec2', region_name=AWS_DEFAULT_REGION)
-
-    # delete default snapshots and images
-    snapshots = ec2_client.describe_snapshots(OwnerIds=['self'])['Snapshots']
-    images = ec2_client.describe_images()['Images']
-    for image in images:
-        ec2_client.deregister_image(ImageId=image.get('ImageId'))
-    for snapshot in snapshots:
-        ec2_client.delete_snapshot(SnapshotId=snapshot.get('SnapshotId'))
-
-    # create infra
-    instance_id = ec2_client.run_instances(ImageId=DEFAULT_AMI_ID, InstanceType=INSTANCE_TYPE_T2_MICRO,
-                                           MaxCount=1, MinCount=1,
-                                           TagSpecifications=[{'ResourceType': 'instance', 'Tags': tags}]
-                                           )['Instances'][0]['InstanceId']
-    image_id = ec2_client.create_image(InstanceId=instance_id, Name=TEST_USER_NAME,
-                                       TagSpecifications=[{'ResourceType': 'image', 'Tags': tags}]).get('ImageId')
-    snapshot_id = (ec2_client.describe_images(ImageIds=[image_id])['Images'][0].get('BlockDeviceMappings')[0]
-                   .get('Ebs').get('SnapshotId'))
-    ec2_client.create_tags(Resources=[snapshot_id], Tags=tags)
-
-    # run zombie_snapshots
-    zombie_snapshots = ZombieSnapshots()
-    response = zombie_snapshots.run()
-    assert len(ec2_client.describe_snapshots(OwnerIds=['self'])['Snapshots']) == 1
-    assert len(response) == 0
+            {'Key': 'DaysCount', 'Value': f'{datetime.utcnow().date()}@7'}]
+    snapshot_id = 'mock_snapshot_id'
+    image_id = 'mock_image_id'
+
+    def mock_create_image(*args, **kwargs):
+        mock_response = {
+            'Images': [
+                {
+                    'ImageId': image_id
+                }
+            ]
+        }
+        return mock_response
+
+    def mock_create_snapshot(*args, **kwargs):
+        return {
+            'Snapshots': [
+                {'SnapshotId': snapshot_id}
+            ]
+        }
+
+    with patch('boto3.client') as mock_client:
+        mock_client.return_value.describe_images.return_value = mock_create_image()
+        mock_client.return_value.describe_snapshots.return_value = mock_create_snapshot()
+
+        # run zombie_snapshots
+        zombie_snapshots = ZombieSnapshots()
+        response = zombie_snapshots.run()
+        ec2_client = boto3.client('ec2', region_name=AWS_DEFAULT_REGION)
+        assert len(ec2_client.describe_snapshots(OwnerIds=['self'])['Snapshots']) == 1
+        assert len(response) == 0