From e47d67bc4fda2455c1791875ef3b256f30e1ad8e Mon Sep 17 00:00:00 2001
From: Anders Ingemann <anders@ingemann.de>
Date: Wed, 8 Apr 2015 21:23:21 +0200
Subject: [PATCH] Launch ec2 test instances inside a VPC

---
 tests/integration/images/ami.py    | 25 ++++++----
 tests/integration/instances/ec2.py | 80 +++++++++++++++++++++++-------
 2 files changed, 76 insertions(+), 29 deletions(-)

diff --git a/tests/integration/images/ami.py b/tests/integration/images/ami.py
index 34714fa..895bf9c 100644
--- a/tests/integration/images/ami.py
+++ b/tests/integration/images/ami.py
@@ -5,31 +5,34 @@ log = logging.getLogger(__name__)
 
 
 def initialize_image(manifest, credentials, bootstrap_info):
-	from boto.ec2 import connect_to_region
-	connection = connect_to_region(bootstrap_info._ec2['region'],
-	                               aws_access_key_id=credentials['access-key'],
-	                               aws_secret_access_key=credentials['secret-key'])
-	image = AmazonMachineImage(manifest, bootstrap_info._ec2['image'], connection)
+	image = AmazonMachineImage(manifest, bootstrap_info._ec2['image'],
+	                           bootstrap_info._ec2['region'], credentials)
 	return image
 
 
 class AmazonMachineImage(Image):
 
-	def __init__(self, manifest, image_id, connection):
+	def __init__(self, manifest, image_id, region, credentials):
 		super(AmazonMachineImage, self).__init__(manifest)
-		self.ami = connection.get_image(image_id)
-		self.connection = connection
+
+		from boto.ec2 import connect_to_region as ec2_connect
+		self.ec2_connection = ec2_connect(region, aws_access_key_id=credentials['access-key'],
+		                                  aws_secret_access_key=credentials['secret-key'])
+		from boto.vpc import connect_to_region as vpc_connect
+		self.vpc_connection = vpc_connect(region, aws_access_key_id=credentials['access-key'],
+		                                  aws_secret_access_key=credentials['secret-key'])
+
+		self.ami = self.ec2_connection.get_image(image_id)
 
 	def destroy(self):
 		log.debug('Deleting AMI')
 		self.ami.deregister()
 		for device, block_device_type in self.ami.block_device_mapping.items():
-			self.connection.delete_snapshot(block_device_type.snapshot_id)
+			self.ec2_connection.delete_snapshot(block_device_type.snapshot_id)
 		del self.ami
 
 	@contextmanager
 	def get_instance(self, instance_type):
 		from ..instances.ec2 import boot_image
-		name = 'bootstrap-vz test instance'
-		with boot_image(name, self.ami, instance_type) as instance:
+		with boot_image(self.ami, instance_type, self.ec2_connection, self.vpc_connection) as instance:
 			yield instance
diff --git a/tests/integration/instances/ec2.py b/tests/integration/instances/ec2.py
index a55e58a..f658856 100644
--- a/tests/integration/instances/ec2.py
+++ b/tests/integration/instances/ec2.py
@@ -5,27 +5,71 @@ log = logging.getLogger(__name__)
 
 
 @contextmanager
-def boot_image(name, image, instance_type):
-	instance = None
-	try:
-		log.debug('Booting ec2 instance')
-		reservation = image.run(instance_type=instance_type)
-		[instance] = reservation.instances
-		instance.add_tag('Name', name)
+def boot_image(image, instance_type, ec2_connection, vpc_connection):
 
-		def instance_running():
-			instance.update()
-			return instance.state == 'running'
-		if not waituntil(instance_running, timeout=120, interval=3):
-			raise EC2InstanceStartupException('Timeout while booting instance')
+	with create_env(ec2_connection, vpc_connection) as boot_env:
 
-		if not waituntil(lambda: instance.get_console_output().output is not None, timeout=600, interval=3):
-			raise EC2InstanceStartupException('Timeout while fetching console output')
+		def waituntil_instance_is(state):
+			def instance_has_state():
+				instance.update()
+				return instance.state == state
+			return waituntil(instance_has_state, timeout=600, interval=3)
 
-		yield instance
-	finally:
-		if instance is not None:
-			instance.terminate()
+		instance = None
+		try:
+			log.debug('Booting ec2 instance')
+			reservation = image.run(instance_type=instance_type,
+			                        subnet_id=boot_env['subnet_id'])
+			[instance] = reservation.instances
+			instance.add_tag('Name', 'bootstrap-vz test instance')
+
+			if not waituntil_instance_is('running'):
+				raise EC2InstanceStartupException('Timeout while booting instance')
+
+			if not waituntil(lambda: instance.get_console_output().output is not None, timeout=600, interval=3):
+				raise EC2InstanceStartupException('Timeout while fetching console output')
+
+			yield instance
+		finally:
+			if instance is not None:
+				log.debug('Terminating ec2 instance')
+				instance.terminate()
+				if not waituntil_instance_is('terminated'):
+					raise EC2InstanceStartupException('Timeout while terminating instance')
+				# wait a little longer, aws can be a little slow sometimes and think the instance is still running
+				import time
+				time.sleep(15)
+
+
+@contextmanager
+def create_env(ec2_connection, vpc_connection):
+
+	vpc_cidr = '10.0.0.0/28'
+	subnet_cidr = '10.0.0.0/28'
+
+	@contextmanager
+	def vpc():
+		log.debug('Creating VPC')
+		vpc = vpc_connection.create_vpc(vpc_cidr)
+		try:
+			yield vpc
+		finally:
+			log.debug('Deleting VPC')
+			vpc_connection.delete_vpc(vpc.id)
+
+	@contextmanager
+	def subnet(vpc):
+		log.debug('Creating subnet')
+		subnet = vpc_connection.create_subnet(vpc.id, subnet_cidr)
+		try:
+			yield subnet
+		finally:
+			log.debug('Deleting subnet')
+			vpc_connection.delete_subnet(subnet.id)
+
+	with vpc() as _vpc:
+		with subnet(_vpc) as _subnet:
+			yield {'subnet_id': _subnet.id}
 
 
 class EC2InstanceStartupException(Exception):