Feature/unit tests refactor (#146)

* Initial add

* Disabled non docdb tests

* Disabled non doc db stack resources

* Fixed up set up test func

* Changed waiter to rds

* Fixed client call

* Enable docdb pull test

* Fixed bugs

* Fixed typo

* Fixed share id

* Fixed docdb pull test

* Disable docdb, enable rds instance

* Fixed rds waiter

* Enabled cloudformation for RDS Instance

* Updated refactored unit tests

* Removed test teardowns

* Moved cleanup before assertion

* Added fix for rds cluster docdb cleanup

* Refactored unit tests

* Re-enabled stack create

* Removed bucket delete

* Moved s3Upload outside script block in pipeline

* Moved s3upload back into script block

Jenkinsfile

@@ -97,9 +97,10 @@ pipeline {
           def safebranch = env.BRANCH_NAME.replace("/", "_")
           def releaseFileName = env.BRANCH_NAME == 'master' ? fileName : fileName.replace('.tar.gz',"-${safebranch}.tar.gz")
           env["SHELVERY_S3_RELEASE"] = "https://${env.SHELVERY_DIST_BUCKET}.s3.amazonaws.com/release/${releaseFileName}"
-
           s3Upload(bucket: env.SHELVERY_DIST_BUCKET, file: "dist/${fileName}", path: "release/${releaseFileName}")
-        }        
+
+        }
+
       }
       post {
         success {

README.md

@@ -325,6 +325,9 @@ backups
                                Note that when copying to a new key, the shelvery requires access to both the new key and the original key.
                                The full KMS key ARN should be supplied.
 
+- `shelvery_reencrypt_kms_key_id`  - when it is required to re-encrypt a RDS instance or cluster snapshot with a different KMS key 
+                                than that of the existing resource. Specify the ARN of the KMS Key.
+
 ### Configuration Priority 0: Sensible defaults
 
 ```text

shelvery/rds_cluster_backup.py

@@ -1,3 +1,4 @@
+from tracemalloc import Snapshot
 import boto3
 
 from shelvery.runtime_config import RuntimeConfig
@@ -303,6 +304,8 @@ def collect_all_snapshots(self, rds_client):
             self.logger.info(f"Collected {len(tmp_snapshots['DBClusterSnapshots'])} manual snapshots. Continuing collection...")
             tmp_snapshots = rds_client.describe_db_cluster_snapshots(SnapshotType='manual', Marker=tmp_snapshots['Marker'])
             all_snapshots.extend(tmp_snapshots['DBClusterSnapshots'])
+
+        all_snapshots = [snapshot for snapshot in all_snapshots if snapshot.get('Engine') != 'docdb']
 
         self.logger.info(f"Collected {len(all_snapshots)} manual snapshots.")
         self.populate_snap_entity_resource(all_snapshots)

shelvery_tests/cleanup_functions.py

@@ -1,117 +1,53 @@
+from shelvery.documentdb_backup import ShelveryDocumentDbBackup
+from shelvery.ebs_backup import ShelveryEBSBackup
+from shelvery.ec2ami_backup import ShelveryEC2AMIBackup
+from shelvery.rds_cluster_backup import ShelveryRDSClusterBackup
+from shelvery.rds_backup import ShelveryRDSBackup
 from shelvery.aws_helper import AwsHelper
+import boto3
+import os
 
 def cleanDocDBSnapshots():
     print("Cleaning up DocDB Snapshots")
-
-    docdbclient = AwsHelper.boto3_client('docdb', region_name='ap-southeast-2')
-
-    snapshots = docdbclient.describe_db_cluster_snapshots(
-        DBClusterIdentifier='shelvery-test-docdb',
-        SnapshotType='Manual'
-    )['DBClusterSnapshots']
-
-    for snapshot in snapshots:
-        snapid = snapshot['DBClusterSnapshotIdentifier']
-
-        try:
-            print(f"Deleting snapshot: {snapid}")
-            docdbclient.delete_db_cluster_snapshot(DBClusterSnapshotIdentifier=snapid)
-        except Exception as e:
-            print(f"Failed to delete {snapid}:{str(e)}")
-
+    backups_engine = ShelveryDocumentDbBackup()
+    backups_engine.clean_backups()
+
 def cleanRdsClusterSnapshots():
     print("Cleaning up RDS Cluster Snapshots")
-    rdsclusterclient = AwsHelper.boto3_client('rds', region_name='ap-southeast-2')
-
-    snapshots =  rdsclusterclient.describe_db_cluster_snapshots(
-        DBClusterIdentifier='shelvery-test-rds-cluster',
-        SnapshotType='Manual'
-    )['DBClusterSnapshots']
-
-    for snapshot in snapshots:
-        snapid = snapshot['DBClusterSnapshotIdentifier']
-
-        try:
-            print(f"Deleting snapshot: {snapid}")
-            rdsclusterclient.delete_db_cluster_snapshot(DBClusterSnapshotIdentifier=snapid)
-        except Exception as e:
-            print(f"Failed to delete {snapid}:{str(e)}")
+    backups_engine = ShelveryRDSClusterBackup()
+    backups_engine.clean_backups()
 
 def cleanRdsSnapshots():
     print("Cleaning up RDS Snapshots")
-    rdsclient = AwsHelper.boto3_client('rds', region_name='ap-southeast-2')
-
-    snapshots = rdsclient.describe_db_snapshots(
-        DBInstanceIdentifier='shelvery-test-rds',
-        SnapshotType='Manual',
-    )['DBSnapshots']
-
-    for snapshot in snapshots:
-        snapid = snapshot['DBSnapshotIdentifier']
-
-        try:
-            print(f"Deleting snapshot: {snapid}")
-            rdsclient.delete_db_snapshot(DBSnapshotIdentifier=snapid)
-        except Exception as e:
-            print(f"Failed to delete {snapid}:{str(e)}")
+    backups_engine = ShelveryRDSBackup()
+    backups_engine.clean_backups()
 
 def cleanEC2Snapshots():
-    #EC2 AMI 
-    ec2client = AwsHelper.boto3_client('ec2', region_name='ap-southeast-2')
-    sts = AwsHelper.boto3_client('sts')
-    id = str(sts.get_caller_identity()['Account'])
-
-    snapshots = ec2client.describe_snapshots(
-        OwnerIds=[id] 
-    )['Snapshots']
-
-    for snapshot in snapshots:
-        snapid = snapshot['SnapshotId']
-        if 'Tags' in snapshot:
-            tags = snapshot['Tags']
-            try:
-                name = [tag['Value'] for tag in tags if tag['Key'] == 'Name'][0]
-                if 'shelvery-test-ec2' in name:
-                    print("Cleaning up EC2 AMI Snapshots")
-                    ami_id = [tag['Value'] for tag in tags if tag['Key'] == 'shelvery:ami_id'][0]
-                    if ami_id != []:
-                        print(f"De-registering image: {ami_id}")
-                        ec2client.deregister_image(ImageId=ami_id)
-                    ec2client.delete_snapshot(SnapshotId=snapid)
-                    print(f'Deleting EC2 snapshot: {snapid}')
-                if 'shelvery-test-ebs' in name:
-                    print("Cleaning up EBS Snapshots")
-                    print(f'Deleting EBS snapshot: {snapid}')
-                    ec2client.delete_snapshot(SnapshotId=snapid)
-            except Exception as e:
-                print(f"Failed to delete {snapid}:{str(e)}")
-
-        else:
-           print(f'Deleting Untagged EC2 Snapshots')
-           if snapshot['VolumeId'] == 'vol-ffffffff' and 'Copied for' in snapshot['Description']:
-
-                search_filter = [{'Name':'block-device-mapping.snapshot-id',
-                                  'Values': [snapid],
-                                  'Name':'tag:ResourceName',
-                                  'Values':['shelvery-test-ec2']
-                                }]
-
-
-
-                ami_id = ec2client.describe_images(
-                    Filters=search_filter
-                )['Images'][0]['ImageId']
-                try:
-                    print(f"De-registering image: {ami_id}")
-                    print(f'Deleting EC2 snapshot: {snapid}')
-                    ec2client.deregister_image(ImageId=ami_id)
-                    ec2client.delete_snapshot(SnapshotId=snapid)
-                except Exception as e:
-                    print(f"Failed to delete {snapid}:{str(e)}")
+    print("Cleaning up EC2 AMI Snapshots")
+    backups_engine = ShelveryEC2AMIBackup()
+    backups_engine.clean_backups()
+
+def cleanEBSSnapshots():
+    print("Cleaning up EBS Snapshots")
+    backups_engine = ShelveryEBSBackup()
+    backups_engine.clean_backups()
+
+def cleanS3Bucket():
+    print("Cleaning S3 Bucket")
+    bucket_name = f"shelvery.data.{AwsHelper.local_account_id()}-ap-southeast-2.base2tools"
+    s3 = boto3.resource('s3')
+    bucket = s3.Bucket(bucket_name)
+
+    # Delete all objects in the bucket
+    for obj in bucket.objects.all():
+        obj.delete()
 
 def cleanupSnapshots():
+    os.environ['shelvery_custom_retention_types'] = 'shortLived:1'
+    os.environ['shelvery_current_retention_type'] = 'shortLived'
     cleanDocDBSnapshots()
     cleanEC2Snapshots()
+    cleanEBSSnapshots()
     cleanRdsClusterSnapshots()
     cleanRdsSnapshots()
 

shelvery_tests/conftest.py

@@ -1,11 +1,10 @@
 import pytest
 import boto3
 import os
-import time
 from shelvery.aws_helper import AwsHelper
-from botocore.exceptions import ClientError, ValidationError, WaiterError
+from botocore.exceptions import ClientError
 
-from shelvery_tests.cleanup_functions import cleanupSnapshots
+from shelvery_tests.cleanup_functions import cleanupSnapshots, cleanS3Bucket
 
 source_account = None
 destination_account = None
@@ -59,6 +58,9 @@ def setup(request):
 
         #Cleanup any existing snapshots after stack is created
         cleanupSnapshots()
+
+        #Cleanup S3 Bucket
+        cleanS3Bucket()
 
         def teardown():
             print ("Initiating Teardown")