From 53c9eb572e84f855a53f8a584c5de05d8cc50416 Mon Sep 17 00:00:00 2001
From: Anders Ingemann <anders@ingemann.de>
Date: Fri, 10 Apr 2015 00:17:00 +0200
Subject: [PATCH] Support testing of EC2 S3 backed instances

---
 .../remote/build_servers/build_server.py      |  3 +++
 tests/integration/providers/ec2/__init__.py   | 25 ++++++++++++++++++
 tests/integration/providers/ec2/images.py     |  8 ++++++
 tests/integration/tools/__init__.py           | 26 ++++++++++++-------
 4 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/bootstrapvz/remote/build_servers/build_server.py b/bootstrapvz/remote/build_servers/build_server.py
index 59c51c1..0cdf6a3 100644
--- a/bootstrapvz/remote/build_servers/build_server.py
+++ b/bootstrapvz/remote/build_servers/build_server.py
@@ -20,4 +20,7 @@ class BuildServer(object):
 			for key in ['access-key', 'secret-key', 'certificate', 'private-key', 'user-id']:
 				if key in self.build_settings['ec2-credentials']:
 					manifest_data['provider']['credentials'][key] = self.build_settings['ec2-credentials'][key]
+		if 's3-region' in self.build_settings and manifest_data['volume']['backing'] == 's3':
+			if 'region' not in manifest_data['image']:
+				manifest_data['image']['region'] = self.build_settings['s3-region']
 		return manifest_data
diff --git a/tests/integration/providers/ec2/__init__.py b/tests/integration/providers/ec2/__init__.py
index 62ef7ce..eda4b1f 100644
--- a/tests/integration/providers/ec2/__init__.py
+++ b/tests/integration/providers/ec2/__init__.py
@@ -4,6 +4,28 @@ import logging
 log = logging.getLogger(__name__)
 
 
+@contextmanager
+def prepare_bootstrap(manifest, build_server):
+	if manifest.volume['backing'] == 's3':
+		credentials = {'access-key': build_server.build_settings['ec2-credentials']['access-key'],
+		               'secret-key': build_server.build_settings['ec2-credentials']['secret-key']}
+		from boto.s3 import connect_to_region as s3_connect
+		s3_connection = s3_connect(manifest.image['region'],
+		                           aws_access_key_id=credentials['access-key'],
+		                           aws_secret_access_key=credentials['secret-key'])
+		log.debug('Creating S3 bucket')
+		bucket = s3_connection.create_bucket(manifest.image['bucket'], location=manifest.image['region'])
+		try:
+			yield
+		finally:
+			log.debug('Deleting S3 bucket')
+			for item in bucket.list():
+				bucket.delete_key(item.key)
+			s3_connection.delete_bucket(manifest.image['bucket'])
+	else:
+			yield
+
+
 @contextmanager
 def boot_image(manifest, build_server, bootstrap_info, instance_type=None):
 
@@ -21,6 +43,9 @@ def boot_image(manifest, build_server, bootstrap_info, instance_type=None):
 	if manifest.volume['backing'] == 'ebs':
 		from images import EBSImage
 		image = EBSImage(bootstrap_info._ec2['image'], ec2_connection)
+	if manifest.volume['backing'] == 's3':
+		from images import S3Image
+		image = S3Image(bootstrap_info._ec2['image'], ec2_connection)
 
 	try:
 		with run_instance(image, instance_type, ec2_connection, vpc_connection) as instance:
diff --git a/tests/integration/providers/ec2/images.py b/tests/integration/providers/ec2/images.py
index b2932fc..6375f8a 100644
--- a/tests/integration/providers/ec2/images.py
+++ b/tests/integration/providers/ec2/images.py
@@ -17,3 +17,11 @@ class EBSImage(AmazonMachineImage):
 		for device, block_device_type in self.ami.block_device_mapping.items():
 			self.ec2_connection.delete_snapshot(block_device_type.snapshot_id)
 		del self.ami
+
+
+class S3Image(AmazonMachineImage):
+
+	def destroy(self):
+		log.debug('Deleting AMI')
+		self.ami.deregister()
+		del self.ami
diff --git a/tests/integration/tools/__init__.py b/tests/integration/tools/__init__.py
index d2c9945..1e82230 100644
--- a/tests/integration/tools/__init__.py
+++ b/tests/integration/tools/__init__.py
@@ -19,17 +19,20 @@ def boot_manifest(manifest_data, boot_vars={}):
 	from bootstrapvz.base.manifest import Manifest
 	manifest = Manifest(data=manifest_data)
 
-	bootstrap_info = None
-	log.info('Connecting to build server')
-	with build_server.connect() as connection:
-		log.info('Building manifest')
-		bootstrap_info = connection.run(manifest)
-
-	log.info('Creating and booting instance')
 	import importlib
 	provider_module = importlib.import_module('tests.integration.providers.' + manifest.provider['name'])
-	with provider_module.boot_image(manifest, build_server, bootstrap_info, **boot_vars) as instance:
-		yield instance
+
+	prepare_bootstrap = getattr(provider_module, 'prepare_bootstrap', noop)
+	with prepare_bootstrap(manifest, build_server):
+		bootstrap_info = None
+		log.info('Connecting to build server')
+		with build_server.connect() as connection:
+			log.info('Building manifest')
+			bootstrap_info = connection.run(manifest)
+
+		log.info('Creating and booting instance')
+		with provider_module.boot_image(manifest, build_server, bootstrap_info, **boot_vars) as instance:
+			yield instance
 
 
 def waituntil(predicate, timeout=5, interval=0.05):
@@ -79,3 +82,8 @@ def read_from_socket(socket_path, termination_string, timeout, read_timeout=0.5)
 				raise SocketReadTimeout(msg)
 		console.close()
 		return output
+
+
+@contextmanager
+def noop(*args, **kwargs):
+	yield