Merge pull request #3 from rightlag/release-0.0.2

Release 0.0.2

Author: rightlag

README.md

@@ -1,6 +1,6 @@
 # timely
 
-timely 0.0.1
+timely 0.0.2
 
 Released: 2-Sep-2015
 
@@ -39,3 +39,15 @@ All commits are tested with [Travis CI](https://travis-ci.org/) and *also* requi
     >>> timely = Timely(verbose=True)
     >>> timely.check()
     Stopping instance: i-6dc5bc92
+
+### Using timezones
+
+By default, `timely` sets the default timezone to `US/Eastern`. However, this can be changed upon instantiation of the `timely` class or via the `set_tz` method. For example:
+
+    >>> from timely import Timely
+    >>> timely = Timely(tz='US/Pacific')
+    >>> timely.set_tz('US/Mountain')
+
+If the timezone specified does not exist, then `UTC` is used.
+
+When assigning the timezone, the `set` method will create a tag `tz` for EC2 containers with the timezone that was specified. Whenever the `check` method is called, the `tz` tag that is assigned to the EC2 container is used to determine whether or not it should be running.

requirements.txt

@@ -1,3 +1,4 @@
 boto==2.38.0
 nose==1.3.7
+pytz==2015.6
 wheel==0.24.0

tests/test_timely.py

@@ -1,5 +1,7 @@
 import boto.ec2
 import datetime
+import pytz
+import sys
 import unittest
 
 from time import sleep
@@ -10,18 +12,19 @@ class TimelyTestCase(unittest.TestCase):
     def setUp(self):
         self.timely = Timely(verbose=True)
         self.conn = boto.ec2.connect_to_region('us-east-1')
-        self.now = datetime.datetime.now()
+        self.now = datetime.datetime.now(tz=pytz.timezone('US/Eastern'))
 
-    def test_times_tag_is_created(self):
+    def test_times_and_tz_tags_are_created(self):
+        """Assert that both the `times` and `tz` tags are created for
+        instances."""
         self.timely.set(weekdays=['*'])
         instances = self.conn.get_only_instances()
         for instance in instances:
-            if instance.state != 'terminated':
-                self.assertIn('times', instance.tags)
-            else:
-                continue
+            self.assertIn('times', instance.tags)
+            self.assertIn('tz', instance.tags)
 
     def test_times_tag_has_length_of_7(self):
+        """Assert that the length of the times tag is 7 elements."""
         self.timely.set(weekdays=['*'])
         instances = self.conn.get_only_instances()
         for instance in instances:
@@ -31,6 +34,7 @@ def test_times_tag_has_length_of_7(self):
             self.assertEqual(len(times), 7)
 
     def test_time_is_set_for_weekday(self):
+        """Assert that a time is set for the current weekday."""
         weekday = self.timely.weekdays[self.now.weekday()]
         self.timely.set(weekdays=[weekday], start_time='9:00 AM',
                         end_time='5:00 PM')
@@ -40,6 +44,9 @@ def test_time_is_set_for_weekday(self):
             self.assertNotEqual(times[self.now.weekday()], str(None))
 
     def test_exception_if_start_time_is_greater_than_equal_to_end_time(self):
+        """If the start time is greater than or equal to the end time
+        a `ValueError` should be raised.
+        """
         with self.assertRaises(ValueError):
             # Greater
             self.timely.set(weekdays=['*'], start_time='9:00 AM',
@@ -49,19 +56,18 @@ def test_exception_if_start_time_is_greater_than_equal_to_end_time(self):
                             end_time='9:00 AM')
 
     def test_unset_method(self):
+        """Assert that the times are set to `None` for all weekdays."""
         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
+        instances = self.conn.get_only_instances()
+        for instance in instances:
+            times = instance.tags['times'].split(';')
             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
+        self.timely.unset(weekdays=['*'])
+        for instance in instances:
+            instance.update()
+            times = instance.tags['times']
+            self.assertEqual(times, ';'.join([str(None)] * 7))
 
     def test_check_method_stops_instance_if_should_not_be_running(self):
         """Check to ensure that an instance is stopped if it SHOULD NOT
@@ -70,36 +76,29 @@ def test_check_method_stops_instance_if_should_not_be_running(self):
         try:
             instance = self.conn.get_only_instances()[0]
             if instance.state == 'stopped':
+                running = False
                 # Start the instance to ensure it is running
+                sys.stdout.write('starting instance: {0}\n'
+                                 .format(instance.id))
                 instance.start()
-                running = False
                 while not running:
-                    try:
-                        instance = self.conn.get_only_instances()[0]
-                        if instance.state == 'running':
-                            running = True
-                        else:
-                            sleep(1)
-                    except IndexError:
-                        pass
+                    instance.update()
+                    if instance.state == 'running':
+                        running = True
+                    else:
+                        sleep(1)
             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
+                instance.update()
+                if instance.state == 'stopped':
+                    stopped = True
+                else:
+                    sleep(1)
             self.assertEqual(instance.state, 'stopped')
         except IndexError:
             pass

timely/timely.py

@@ -1,12 +1,14 @@
 import boto.ec2
 import collections
+import pytz
 import sys
 
 from datetime import datetime
 
 
 class Timely(object):
-    def __init__(self, region='us-east-1', iso=False, verbose=False):
+    def __init__(self, region='us-east-1', iso=False, verbose=False,
+                 tz='US/Eastern'):
         self.conn = boto.ec2.connect_to_region(region)
         self.iso = iso
         # Accommodate for ISO weekday - remove first element if `iso` is False
@@ -27,6 +29,7 @@ def __init__(self, region='us-east-1', iso=False, verbose=False):
             weekdays.pop(0)
             self.weekdays = weekdays
         self.verbose = verbose
+        self.set_tz(tz)
 
     def use_verbose(self):
         self.verbose = True
@@ -37,6 +40,15 @@ def use_iso(self):
     def set_region(self, region):
         self.conn = boto.ec2.connect_to_region(region)
 
+    def set_tz(self, tz):
+        try:
+            self.tz = pytz.timezone(tz)
+        except pytz.exceptions.UnknownTimeZoneError:
+            self.tz = pytz.utc
+
+    def _verbose_message(self, action, instance):
+        sys.stdout.write('{0} instance: {1}\n'.format(action, instance.id))
+
     def all(self, instance_ids=None):
         """Read weekday run times for all or specific EC2 instances.
 
@@ -60,8 +72,12 @@ def all(self, instance_ids=None):
                         end_time = datetime.strptime(end_time, '%H:%M')
                     except ValueError:
                         continue
-                    start_time = start_time.strftime('%H:%M')
-                    end_time = end_time.strftime('%H:%M')
+                    tz = instance.tags.get('tz', pytz.utc)
+                    tz = pytz.timezone(tz)
+                    start_time = tz.localize(start_time)
+                    end_time = tz.localize(end_time)
+                    start_time = start_time.strftime('%I:%M %p %Z%z')
+                    end_time = end_time.strftime('%I:%M %p %Z%z')
                     weekday = (self.weekdays[i + 1]
                                if self.iso else self.weekdays[i])
                     data[instance.id].append(
@@ -105,10 +121,14 @@ def set(self, instance_ids=None, weekdays=None, start_time=None,
                 continue
             times = instance.tags.get('times')
             if not times:
-                # No `times` tag - set default
+                # No `times` tag - set defaults
                 times = ';'.join([str(None)] * 7)
+                tags = {
+                    'times': times,
+                    'tz': self.tz,
+                }
                 try:
-                    instance.add_tag('times', times)
+                    instance.add_tags(tags)
                 except self.conn.ResponseError, e:
                     raise e
             times = times.split(';')
@@ -135,9 +155,13 @@ def set(self, instance_ids=None, weekdays=None, start_time=None,
                 # Remove first element `None` from `times` object
                 times.pop(0)
             times = ';'.join([str(time) for time in times])
+            tags = {
+                'times': times,
+                'tz': self.tz,
+            }
             try:
                 # Overwrite existing `times` tag with new value
-                instance.add_tag('times', times)
+                instance.add_tags(tags)
             except self.conn.ResponseError, e:
                 raise e
 
@@ -180,7 +204,13 @@ def check(self, instance_ids=None):
         """
         instances = self.conn.get_only_instances(instance_ids=instance_ids)
         for instance in instances:
-            now = datetime.now()
+            tz = instance.tags.get('tz')
+            if not tz:
+                # Needs to be a timezone to compare current time to
+                # `start_time` and `end_time` - otherwise continue
+                continue
+            tz = pytz.timezone(tz)
+            now = datetime.now(tz=tz)
             if self.iso:
                 weekday = now.isoweekday()
             else:
@@ -206,34 +236,25 @@ def check(self, instance_ids=None):
                                                     day=now.day)
                     end_time = end_time.replace(year=now.year,
                                                 month=now.month, day=now.day)
+                    # http://www.saltycrane.com/blog/2009/05/converting-time-zones-datetime-objects-python/#add-timezone-localize
+                    start_time = tz.localize(start_time)
+                    end_time = tz.localize(end_time)
                     if start_time <= now <= end_time:
                         if instance.state == 'stopped':
                             if self.verbose:
-                                sys.stdout.write(
-                                    'Starting instance: {0}\r\n'.format(
-                                        instance.id
-                                    )
-                                )
+                                self._verbose_message('starting', instance)
                             instance.start()
                     else:
                         if instance.state == 'running':
                             if self.verbose:
-                                sys.stdout.write(
-                                    'Stopping instance: {0}\r\n'.format(
-                                        instance.id
-                                    )
-                                )
+                                self._verbose_message('stopping', instance)
                             instance.stop()
                 else:
                     # If the time is `None` check to see if the instance is
                     # running - if it is, then stop it by default
                     if instance.state == 'running':
                         if self.verbose:
-                            sys.stdout.write(
-                                'Stopping instance: {0}\r\n'.format(
-                                    instance.id
-                                )
-                            )
+                            self._verbose_message('stopping', instance)
                         instance.stop()
 
     def __str__(self):