From 9f19662d1c3227638c7f06e72ed28a0ec59802a2 Mon Sep 17 00:00:00 2001
From: Steve Pulec <spulec@gmail.com>
Date: Fri, 17 May 2013 19:35:53 -0400
Subject: [PATCH] allow passing user data to run_instances

---
 moto/ec2/models.py               | 49 ++++++++++++++++++++------------
 moto/ec2/responses/instances.py  |  3 +-
 tests/test_ec2/test_instances.py | 15 ++++++++++
 3 files changed, 48 insertions(+), 19 deletions(-)

diff --git a/moto/ec2/models.py b/moto/ec2/models.py
index c254e1716..f69919aec 100644
--- a/moto/ec2/models.py
+++ b/moto/ec2/models.py
@@ -16,10 +16,29 @@ from .utils import (
 
 
 class Instance(BotoInstance):
-    def __init__(self):
-        self._state_name = None
-        self._state_code = None
+    def __init__(self, image_id, user_data):
         super(Instance, self).__init__()
+        self.id = random_instance_id()
+        self.image_id = image_id
+        self._state_name = "pending"
+        self._state_code = 0
+        self.user_data = user_data
+
+    def start(self):
+        self._state_name = "pending"
+        self._state_code = 0
+
+    def stop(self):
+        self._state_name = "stopping"
+        self._state_code = 64
+
+    def terminate(self):
+        self._state_name = "shutting-down"
+        self._state_code = 32
+
+    def reboot(self):
+        self._state_name = "pending"
+        self._state_code = 0
 
 
 class InstanceBackend(object):
@@ -33,15 +52,14 @@ class InstanceBackend(object):
             if instance.id == instance_id:
                 return instance
 
-    def add_instances(self, image_id, count):
+    def add_instances(self, image_id, count, user_data):
         new_reservation = Reservation()
         new_reservation.id = random_reservation_id()
         for index in range(count):
-            new_instance = Instance()
-            new_instance.id = random_instance_id()
-            new_instance.image_id = image_id
-            new_instance._state_name = "pending"
-            new_instance._state_code = 0
+            new_instance = Instance(
+                image_id,
+                user_data,
+            )
             new_reservation.instances.append(new_instance)
         self.reservations[new_reservation.id] = new_reservation
         return new_reservation
@@ -50,8 +68,7 @@ class InstanceBackend(object):
         started_instances = []
         for instance in self.all_instances():
             if instance.id in instance_ids:
-                instance._state_name = "pending"
-                instance._state_code = 0
+                instance.start()
                 started_instances.append(instance)
 
         return started_instances
@@ -60,8 +77,7 @@ class InstanceBackend(object):
         stopped_instances = []
         for instance in self.all_instances():
             if instance.id in instance_ids:
-                instance._state_name = "stopping"
-                instance._state_code = 64
+                instance.stop()
                 stopped_instances.append(instance)
 
         return stopped_instances
@@ -70,8 +86,7 @@ class InstanceBackend(object):
         terminated_instances = []
         for instance in self.all_instances():
             if instance.id in instance_ids:
-                instance._state_name = "shutting-down"
-                instance._state_code = 32
+                instance.terminate()
                 terminated_instances.append(instance)
 
         return terminated_instances
@@ -80,9 +95,7 @@ class InstanceBackend(object):
         rebooted_instances = []
         for instance in self.all_instances():
             if instance.id in instance_ids:
-                # TODO double check instances go to pending when reboot
-                instance._state_name = "pending"
-                instance._state_code = 0
+                instance.reboot()
                 rebooted_instances.append(instance)
 
         return rebooted_instances
diff --git a/moto/ec2/responses/instances.py b/moto/ec2/responses/instances.py
index 7170a0928..d752bf0f6 100644
--- a/moto/ec2/responses/instances.py
+++ b/moto/ec2/responses/instances.py
@@ -13,7 +13,8 @@ class InstanceResponse(object):
     def run_instances(self):
         min_count = int(self.querystring.get('MinCount', ['1'])[0])
         image_id = self.querystring.get('ImageId')[0]
-        new_reservation = ec2_backend.add_instances(image_id, min_count)
+        user_data = self.querystring.get('UserData')
+        new_reservation = ec2_backend.add_instances(image_id, min_count, user_data)
         template = Template(EC2_RUN_INSTANCES)
         return template.render(reservation=new_reservation)
 
diff --git a/tests/test_ec2/test_instances.py b/tests/test_ec2/test_instances.py
index 618b4d1ff..47e8b4b16 100644
--- a/tests/test_ec2/test_instances.py
+++ b/tests/test_ec2/test_instances.py
@@ -1,3 +1,5 @@
+import base64
+
 import boto
 from boto.ec2.instance import Reservation, InstanceAttribute
 import sure  # flake8: noqa
@@ -98,3 +100,16 @@ def test_instance_attribute_user_data():
     instance_attribute = instance.get_attribute("userData")
     instance_attribute.should.be.a(InstanceAttribute)
     instance_attribute.get("userData").should.equal("this is my user data")
+
+
+@mock_ec2
+def test_user_data_with_run_instance():
+    user_data = "some user data"
+    conn = boto.connect_ec2('the_key', 'the_secret')
+    reservation = conn.run_instances('ami-1234abcd', user_data=user_data)
+    instance = reservation.instances[0]
+
+    instance_attribute = instance.get_attribute("userData")
+    instance_attribute.should.be.a(InstanceAttribute)
+    decoded_user_data = base64.decodestring(instance_attribute.get("userData"))
+    decoded_user_data.should.equal("some user data")