Add unit test to spark_ec2 script

.gitignore

@@ -3,6 +3,7 @@
 *.ipr
 *.iml
 *.iws
+*.pyc
 .idea/
 sbt/*.jar
 .settings

LICENSE

@@ -236,6 +236,54 @@ WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
 IN THE SOFTWARE.
 
+=======================================================================
+For the moto library (ec2/third_party/moto*.zip)
+=======================================================================
+
+Copyright 2012 Steve Pulec
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+=======================================================================
+For the mock library (ec2/third_party/mock*.zip)
+=======================================================================
+
+Copyright (c) 2003-2012, Michael Foord
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+    * Redistributions of source code must retain the above copyright
+      notice, this list of conditions and the following disclaimer.
+
+    * Redistributions in binary form must reproduce the above
+      copyright notice, this list of conditions and the following
+      disclaimer in the documentation and/or other materials provided
+      with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 ========================================================================
 For CloudPickle (pyspark/cloudpickle.py):

ec2/spark_ec2.py

@@ -37,6 +37,8 @@
 from boto.ec2.blockdevicemapping import BlockDeviceMapping, EBSBlockDeviceType
 from boto import ec2
 
+AWS_EVENTUAL_CONSISTENCY = 30
+
 class UsageError(Exception):
   pass
 
@@ -223,7 +225,7 @@ def launch_cluster(conn, opts, cluster_name):
     sys.exit(1)
   if opts.key_pair is None:
     print >> stderr, "ERROR: Must provide a key pair name (-k) to use on instances."
-    sys.exit(1)    
+    sys.exit(1)
   print "Setting up security groups..."
   master_group = get_or_make_group(conn, cluster_name + "-master")
   slave_group = get_or_make_group(conn, cluster_name + "-slaves")
@@ -459,7 +461,6 @@ def setup_spark_cluster(master, opts):
   if opts.ganglia:
     print "Ganglia started at http://%s:5080/ganglia" % master
 
-
 # Wait for a whole cluster (masters, slaves and ZooKeeper) to start up
 def wait_for_cluster(conn, wait_secs, master_nodes, slave_nodes):
   print "Waiting for instances to start up..."
@@ -665,6 +666,59 @@ def get_partition(total, num_partitions, current_partitions):
   return num_slaves_this_zone
 
 
+def destroy_cluster(conn, opts, cluster_name):
+  (master_nodes, slave_nodes) = get_existing_cluster(
+      conn, opts, cluster_name, die_on_error=False)
+  print "Terminating master..."
+  for inst in master_nodes:
+    inst.terminate()
+  print "Terminating slaves..."
+  for inst in slave_nodes:
+    inst.terminate()
+
+  # Delete security groups as well
+  if opts.delete_groups:
+    print "Deleting security groups (this will take some time)..."
+    group_names = [cluster_name + "-master", cluster_name + "-slaves"]
+
+    attempt = 1;
+    while attempt <= 3:
+      print "Attempt %d" % attempt
+      groups = [g for g in conn.get_all_security_groups() if g.name in group_names]
+      success = True
+      # Delete individual rules in all groups before deleting groups to
+      # remove dependencies between them
+      for group in groups:
+        print "Deleting rules in security group " + group.name
+        for rule in group.rules:
+          for grant in rule.grants:
+              success &= group.revoke(ip_protocol=rule.ip_protocol,
+                       from_port=rule.from_port,
+                       to_port=rule.to_port,
+                       src_group=grant)
+
+      # Sleep for AWS eventual-consistency to catch up, and for instances
+      # to terminate
+      time.sleep(AWS_EVENTUAL_CONSISTENCY)  # Yes, it does have to be this long :-(
+      for group in groups:
+        try:
+          conn.delete_security_group(group.name)
+          print "Deleted security group " + group.name
+        except boto.exception.EC2ResponseError:
+          success = False;
+          print "Failed to delete security group " + group.name
+
+      # Unfortunately, group.revoke() returns True even if a rule was not
+      # deleted, so this needs to be rerun if something fails
+      if success: break;
+
+      attempt += 1
+
+    if not success:
+      print "Failed to delete all security groups after 3 tries."
+      print "Try re-running in a few minutes."
+
+
 def real_main():
   (opts, action, cluster_name) = parse_args()
   try:
@@ -695,56 +749,7 @@ def real_main():
         cluster_name + "?\nALL DATA ON ALL NODES WILL BE LOST!!\n" +
         "Destroy cluster " + cluster_name + " (y/N): ")
     if response == "y":
-      (master_nodes, slave_nodes) = get_existing_cluster(
-          conn, opts, cluster_name, die_on_error=False)
-      print "Terminating master..."
-      for inst in master_nodes:
-        inst.terminate()
-      print "Terminating slaves..."
-      for inst in slave_nodes:
-        inst.terminate()
-
-      # Delete security groups as well
-      if opts.delete_groups:
-        print "Deleting security groups (this will take some time)..."
-        group_names = [cluster_name + "-master", cluster_name + "-slaves"]
-
-        attempt = 1;
-        while attempt <= 3:
-          print "Attempt %d" % attempt
-          groups = [g for g in conn.get_all_security_groups() if g.name in group_names]
-          success = True
-          # Delete individual rules in all groups before deleting groups to
-          # remove dependencies between them
-          for group in groups:
-            print "Deleting rules in security group " + group.name
-            for rule in group.rules:
-              for grant in rule.grants:
-                  success &= group.revoke(ip_protocol=rule.ip_protocol,
-                           from_port=rule.from_port,
-                           to_port=rule.to_port,
-                           src_group=grant)
-
-          # Sleep for AWS eventual-consistency to catch up, and for instances
-          # to terminate
-          time.sleep(30)  # Yes, it does have to be this long :-(
-          for group in groups:
-            try:
-              conn.delete_security_group(group.name)
-              print "Deleted security group " + group.name
-            except boto.exception.EC2ResponseError:
-              success = False;
-              print "Failed to delete security group " + group.name
-
-          # Unfortunately, group.revoke() returns True even if a rule was not
-          # deleted, so this needs to be rerun if something fails
-          if success: break;
-
-          attempt += 1
-
-        if not success:
-          print "Failed to delete all security groups after 3 tries."
-          print "Try re-running in a few minutes."
+      destroy_cluster()
 
   elif action == "login":
     (master_nodes, slave_nodes) = get_existing_cluster(

ec2/tests.py

@@ -0,0 +1,25 @@
+import unittest
+
+from boto import ec2
+import mock
+import moto
+
+import spark_ec2
+
+
+class CommandTests(unittest.TestCase):
+
+    @moto.mock_ec2
+    def test_destroy(self):
+        spark_ec2.AWS_EVENTUAL_CONSISTENCY = 1
+        opts = mock.MagicMock(name='opts')
+        opts.region = "us-east-1"
+        conn = ec2.connect_to_region(opts.region)
+        cluster_name = "cluster_name"
+        try:
+            spark_ec2.destroy_cluster(conn, opts, cluster_name)
+        except:
+            self.fail("destroy_cluster raised unexpected exception")
+
+if __name__ == '__main__':
+    unittest.main()