Add timezone option to config

Using dateutil.tz to link string names to tzinfo objects,
the create_date can now generate using a named timezone
rather than datetime.now().

Change-Id: I9f151cb9e11da3d68be63d7141f60e7eccb9812c
Fixes: #425

1 files changed, 67 insertions, 0 deletions

diff --git a/tests/test_script_production.py b/tests/test_script_production.py
index 3703014..3364d56 100644
--- a/tests/test_script_production.py
+++ b/tests/test_script_production.py
@@ -18,6 +18,7 @@ import sqlalchemy as sa
 from sqlalchemy.engine.reflection import Inspector
 from alembic.util import CommandError
 import re
+from dateutil import tz
 
 env, abc, def_ = None, None, None
 
@@ -153,6 +154,72 @@ class ScriptNamingTest(TestBase):
                 "message_2012_7_25_15_8_5.py" % _get_staging_directory())
         )
 
+    def _test_tz(self, timezone_arg, given, expected):
+        script = ScriptDirectory(
+            _get_staging_directory(),
+            file_template="%(rev)s_%(slug)s_"
+            "%(year)s_%(month)s_"
+            "%(day)s_%(hour)s_"
+            "%(minute)s_%(second)s",
+            timezone=timezone_arg
+        )
+
+        with mock.patch(
+                "alembic.script.base.datetime",
+                mock.Mock(
+                    datetime=mock.Mock(
+                        utcnow=lambda: given,
+                        now=lambda: given
+                    )
+                )
+        ):
+            create_date = script._generate_create_date()
+        eq_(
+            create_date,
+            expected
+        )
+
+    def test_custom_tz(self):
+        self._test_tz(
+            'EST5EDT',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT'))
+        )
+
+    def test_custom_tz_lowercase(self):
+        self._test_tz(
+            'est5edt',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 11, 8, 5, tzinfo=tz.gettz('EST5EDT'))
+        )
+
+    def test_custom_tz_utc(self):
+        self._test_tz(
+            'utc',
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(
+                2012, 7, 25, 15, 8, 5, tzinfo=tz.gettz('UTC'))
+        )
+
+    def test_default_tz(self):
+        self._test_tz(
+            None,
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(2012, 7, 25, 15, 8, 5)
+        )
+
+    def test_tz_cant_locate(self):
+        assert_raises_message(
+            CommandError,
+            "Can't locate timezone: fake",
+            self._test_tz,
+            "fake",
+            datetime.datetime(2012, 7, 25, 15, 8, 5),
+            datetime.datetime(2012, 7, 25, 15, 8, 5)
+        )
+
 
 class RevisionCommandTest(TestBase):
     def setUp(self):