1 files changed, 126 insertions, 0 deletions
diff --git a/tests/cloud_tests/platforms/ec2/instance.py b/tests/cloud_tests/platforms/ec2/instance.py
new file mode 100644
index 00000000..4ba737ab
--- /dev/null
+++ b/tests/cloud_tests/platforms/ec2/instance.py
@@ -0,0 +1,126 @@
+# This file is part of cloud-init. See LICENSE file for license information.
+
+"""Base EC2 instance."""
+import os
+
+import botocore
+
+from ..instances import Instance
+from tests.cloud_tests import LOG, util
+
+
+class EC2Instance(Instance):
+    """EC2 backed instance."""
+
+    platform_name = "ec2"
+    _ssh_client = None
+
+    def __init__(self, platform, properties, config, features,
+                 image_ami, user_data=None):
+        """Set up instance.
+
+        @param platform: platform object
+        @param properties: dictionary of properties
+        @param config: dictionary of configuration values
+        @param features: dictionary of supported feature flags
+        @param image_ami: AWS AMI ID for image to use
+        @param user_data: test user-data to pass to instance
+        """
+        super(EC2Instance, self).__init__(
+            platform, image_ami, properties, config, features)
+
+        self.image_ami = image_ami
+        self.instance = None
+        self.user_data = user_data
+        self.ssh_ip = None
+        self.ssh_port = 22
+        self.ssh_key_file = os.path.join(
+            platform.config['data_dir'], platform.config['private_key'])
+        self.ssh_pubkey_file = os.path.join(
+            platform.config['data_dir'], platform.config['public_key'])
+
+    def console_log(self):
+        """Collect console log from instance.
+
+        The console log is buffered and not always present, therefore
+        may return empty string.
+        """
+        try:
+            return self.instance.console_output()['Output'].encode()
+        except KeyError:
+            return b''
+
+    def destroy(self):
+        """Clean up instance."""
+        if self.instance:
+            LOG.debug('destroying instance %s', self.instance.id)
+            self.instance.terminate()
+            self.instance.wait_until_terminated()
+
+        self._ssh_close()
+
+        super(EC2Instance, self).destroy()
+
+    def _execute(self, command, stdin=None, env=None):
+        """Execute command on instance."""
+        env_args = []
+        if env:
+            env_args = ['env'] + ["%s=%s" for k, v in env.items()]
+
+        return self._ssh(['sudo'] + env_args + list(command), stdin=stdin)
+
+    def start(self, wait=True, wait_for_cloud_init=False):
+        """Start instance on EC2 with the platfrom's VPC."""
+        if self.instance:
+            if self.instance.state['Name'] == 'running':
+                return
+
+            LOG.debug('starting instance %s', self.instance.id)
+            self.instance.start()
+        else:
+            LOG.debug('launching instance')
+
+            args = {
+                'ImageId': self.image_ami,
+                'InstanceType': self.platform.instance_type,
+                'KeyName': self.platform.key_name,
+                'MaxCount': 1,
+                'MinCount': 1,
+                'SecurityGroupIds': [self.platform.security_group.id],
+                'SubnetId': self.platform.subnet.id,
+                'TagSpecifications': [{
+                    'ResourceType': 'instance',
+                    'Tags': [{
+                        'Key': 'Name', 'Value': self.platform.tag
+                    }]
+                }],
+            }
+
+            if self.user_data:
+                args['UserData'] = self.user_data
+
+            try:
+                instances = self.platform.ec2_resource.create_instances(**args)
+            except botocore.exceptions.ClientError as error:
+                error_msg = error.response['Error']['Message']
+                raise util.PlatformError('start', error_msg)
+
+            self.instance = instances[0]
+
+        LOG.debug('instance id: %s', self.instance.id)
+        if wait:
+            self.instance.wait_until_running()
+            self.instance.reload()
+            self.ssh_ip = self.instance.public_ip_address
+            self._wait_for_system(wait_for_cloud_init)
+
+    def shutdown(self, wait=True):
+        """Shutdown instance."""
+        LOG.debug('stopping instance %s', self.instance.id)
+        self.instance.stop()
+
+        if wait:
+            self.instance.wait_until_stopped()
+            self.instance.reload()
+
+# vi: ts=4 expandtab