Launch ec2 test instances inside a VPC

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
+	image = AmazonMachineImage(manifest, bootstrap_info._ec2['image'],
-	connection = connect_to_region(bootstrap_info._ec2['region'],
+	                           bootstrap_info._ec2['region'], credentials)
-	                               aws_access_key_id=credentials['access-key'],
-	                               aws_secret_access_key=credentials['secret-key'])
-	image = AmazonMachineImage(manifest, bootstrap_info._ec2['image'], connection)
 	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(self.ami, instance_type, self.ec2_connection, self.vpc_connection) as instance:
-		with boot_image(name, self.ami, instance_type) as instance:
 			yield instance

tests/integration/instances/ec2.py

@@ -5,18 +5,25 @@ log = logging.getLogger(__name__)
 
 
 @contextmanager
-def boot_image(name, image, instance_type):
+def boot_image(image, instance_type, ec2_connection, vpc_connection):
+
+	with create_env(ec2_connection, vpc_connection) as boot_env:
+
+		def waituntil_instance_is(state):
+			def instance_has_state():
+				instance.update()
+				return instance.state == state
+			return waituntil(instance_has_state, timeout=600, interval=3)
+
 		instance = None
 		try:
 			log.debug('Booting ec2 instance')
-		reservation = image.run(instance_type=instance_type)
+			reservation = image.run(instance_type=instance_type,
+			                        subnet_id=boot_env['subnet_id'])
 			[instance] = reservation.instances
-		instance.add_tag('Name', name)
+			instance.add_tag('Name', 'bootstrap-vz test instance')
 
-		def instance_running():
+			if not waituntil_instance_is('running'):
-			instance.update()
-			return instance.state == 'running'
-		if not waituntil(instance_running, timeout=120, interval=3):
 				raise EC2InstanceStartupException('Timeout while booting instance')
 
 			if not waituntil(lambda: instance.get_console_output().output is not None, timeout=600, interval=3):
@@ -25,7 +32,44 @@ def boot_image(name, image, instance_type):
 			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):