#!/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() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment