Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
#!/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
You can’t perform that action at this time.