empty_buckets.py

#!/usr/bin/env python3
import subprocess
import json

# bucket_template = 'fugue-e2e-s3-%s-logging'
def empty_all_buckets(bucket_template):
    regions = get_regions()

    for r in regions:
        bucket_name =  bucket_template % (r)
        if (verify_bucket_exists(bucket_name, r)):
            print('Bucket exists. Emptying.')
            empty_bucket(bucket_name, r)


def get_regions():
    regions = []
    call = "aws ec2 describe-regions --query 'Regions[].{Region:RegionName}'"
    try:
        output = subprocess.check_output([call, ""], shell=True)

        formatted_output = json.loads(output.decode('utf-8'))

        for r in formatted_output:
            if 'Region' in r and r['Region'] is not None:
                regions.append(r['Region'])
    except Exception as e:
        return 'Error: unable to get regions: %s' % e

    return regions


def verify_bucket_exists(bucket, region):
    call = "aws --region %s s3api head-bucket --bucket %s" % (region, bucket) 

    try:
        output = subprocess.check_output([call, ""], shell=True)
        if output:
            raise Exception('does not exist')
        else:
            return True
    # if the bucket doesn't exist, there will be an error
    except Exception as e:
        return 'Error: %s %s' % (bucket, e)


def empty_bucket(bucket, region):
    call = "aws --region %s s3 rm s3://%s --recursive" % (region, bucket)
    try:
        output = subprocess.check_output([call, ""], shell=True)
        formatted_output = output.decode('utf-8')

        return formatted_output
    except Exception as e:
        error_message = 'Error: logging bucket in %s %s' % (region, e)

        return error_message
      
# ------

import unittest
from unittest import TestCase
from unittest.mock import patch
import empty_buckets as eb
from subprocess import CalledProcessError

# pulled directly from AWS CLI call
expected_region_val = '[{"Region": "ap-south-1"},{"Region":"eu-west-3"}]'.encode('utf-8')
malformed_region_val = '[{"Region": "ap-south-1"},{"Region":null}]'.encode('utf-8')

class TestGetRegions(TestCase):
    @patch('subprocess.check_output')
    def test_get_regions_works_as_expected(self, mock_subprocess):
        mock_subprocess.return_value=expected_region_val

        regions = eb.get_regions()
        self.assertEqual(regions, ['ap-south-1', 'eu-west-3'])

    @patch('subprocess.check_output')
    def test_get_regions_works_malformed_value(self, mock_subprocess):
        mock_subprocess.return_value=malformed_region_val

        regions = eb.get_regions()
        self.assertEqual(regions, ['ap-south-1'])
        

class TestVerifyBucketExists(TestCase):
    @patch('subprocess.check_output')
    def test_verify_bucket_exists_true(self, mock_subprocess):
        mock_subprocess.return_value=''

        result = eb.verify_bucket_exists('sample_bucket_doesnt_matter', 'us-east-1')
        self.assertEqual(result, True)

    @patch('subprocess.check_output')
    def test_verify_bucket_exists_false(self, mock_subprocess):
        mock_subprocess.return_value='An error occurred (404) when calling the HeadBucket operation: Not Found'

        result = eb.verify_bucket_exists('sample_bucket_doesnt_matter', 'us-east-1')
        self.assertEqual(result, 'Error: sample_bucket_doesnt_matter does not exist')


expected_output = 'delete: s3://mybucket/test1.txt \ndelete: s3://mybucket/test2.txt'

class TestEmptyBucket(TestCase):
    @patch('subprocess.check_output')
    def test_empty_bucket_with_output(self, mock_subprocess):
        mock_subprocess.return_value=expected_output.encode('utf-8')

        result = eb.empty_bucket('sample_bucket_doesnt_matter', 'us-east-1')
        self.assertEqual(result, expected_output)

    @patch('subprocess.check_output')
    def test_empty_bucket_with_error(self, mock_subprocess):
        mock_subprocess.side_effect = CalledProcessError

        self.assertRaises(Exception, eb.empty_bucket('sample_bucket_doesnt_matter', 'us-east-1'))

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