Merge pull request #2 from rightlag/develop

Develop

Author: rightlag

tests/test_timely.py

@@ -2,6 +2,7 @@
 import datetime
 import unittest
 
+from time import sleep
 from timely import Timely
 
 
@@ -47,5 +48,48 @@ def test_exception_if_start_time_is_greater_than_equal_to_end_time(self):
             self.timely.set(weekdays=['*'], start_time='9:00 AM',
                             end_time='9:00 AM')
 
+    def test_unset_method(self):
+        self.timely.set(weekdays=['*'], start_time='9:00 AM',
+                        end_time='5:00 PM')
+        try:
+            instance = self.conn.get_only_instances()[0]
+            times = self.timely.all()[instance.id]
+            # First - set times for all days of the week
+            self.assertEqual(len(times), 7)
+            # Second - unset times for all days of the week
+            self.timely.unset(weekdays=['*'])
+            times = self.timely.all()[instance.id]
+            self.assertEqual(len(times), 0)
+        except IndexError:
+            pass
+
+    def test_check_method_stops_instance_if_should_not_be_running(self):
+        try:
+            instance = self.conn.get_only_instances()[0]
+            if instance.state == 'stopped':
+                # Start the instance to ensure it is running
+                instance.start()
+            weekday = self.timely.weekdays[self.now.weekday()]
+            # Automatically sets `start_time` and `end_time` to `None`
+            self.timely.set(weekdays=[weekday])
+            # Ensure that the instance is being stopped
+            self.timely.check()
+            stopped = False
+            instance = None
+            while not stopped:
+                try:
+                    # Need to remake a connection to AWS to get the updated
+                    # instance status
+                    instance = self.conn.get_only_instances()[0]
+                    if instance.state == 'stopped':
+                        stopped = True
+                    else:
+                        sleep(1)
+                except IndexError:
+                    pass
+            self.assertEqual(instance.state, 'stopped')
+        except IndexError:
+            pass
+
     def tearDown(self):
         del self.timely

timely/timely.py

@@ -1,4 +1,5 @@
 import boto.ec2
+import collections
 import sys
 
 from datetime import datetime
@@ -30,6 +31,12 @@ def __init__(self, region='us-east-1', iso=False, verbose=False):
     def use_verbose(self):
         self.verbose = True
 
+    def use_iso(self):
+        self.iso = True
+
+    def set_region(self, region):
+        self.conn = boto.ec2.connect_to_region(region)
+
     def all(self, instance_ids=None):
         """Read weekday run times for all or specific EC2 instances.
 
@@ -38,6 +45,8 @@ def all(self, instance_ids=None):
                 IDs
         """
         data = {}
+        Time = collections.namedtuple('Time',
+                                      ['weekday', 'start_time', 'end_time'])
         instances = self.conn.get_only_instances(instance_ids=instance_ids)
         for instance in instances:
             times = instance.tags.get('times')
@@ -55,7 +64,9 @@ def all(self, instance_ids=None):
                     end_time = end_time.strftime('%H:%M')
                     weekday = (self.weekdays[i + 1]
                                if self.iso else self.weekdays[i])
-                    data[instance.id].append((weekday, start_time, end_time,))
+                    data[instance.id].append(
+                        Time(weekday, start_time, end_time)
+                    )
         return data
 
     def set(self, instance_ids=None, weekdays=None, start_time=None,
@@ -75,7 +86,7 @@ def set(self, instance_ids=None, weekdays=None, start_time=None,
             start_time = datetime.strptime(start_time, '%I:%M %p')
             end_time = datetime.strptime(end_time, '%I:%M %p')
             if start_time >= end_time:
-                raise ValueError('start time can\'t be greater than end time')
+                raise ValueError('Start time can\'t be greater than end time')
             start_time = start_time.strftime('%H:%M')
             end_time = end_time.strftime('%H:%M')
             updated = '{0}-{1}'.format(start_time, end_time)
@@ -130,6 +141,28 @@ def set(self, instance_ids=None, weekdays=None, start_time=None,
             except self.conn.ResponseError, e:
                 raise e
 
+    def unset(self, instance_ids=None, weekdays=None):
+        """Unset instance times for specific weekdays or all weekdays."""
+        # integer representation of `weekdays`
+        if weekdays == ['*']:
+            # All 7 days
+            weekdays = range(len(self.weekdays))
+        else:
+            weekdays = [self.weekdays.index(weekday) for weekday in weekdays]
+        instances = self.conn.get_only_instances(instance_ids=instance_ids)
+        for instance in instances:
+            times = instance.tags.get('times')
+            if times:
+                times = times.split(';')
+                for weekday in weekdays:
+                    times[weekday] = None
+                times = ';'.join([str(time) for time in times])
+                try:
+                    # Overwrite existing `times` tag with new value
+                    instance.add_tag('times', times)
+                except self.conn.ResponseError, e:
+                    raise e
+
     def check(self, instance_ids=None):
         """Check the state of instances and either start or stop them
         based on the current time restrictions set for the current day.