diff --git a/autopush/db.py b/autopush/db.py index 0a093f84..641295c9 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -129,10 +129,10 @@ def make_rotating_tablename(prefix, delta=0, date=None): def create_rotating_message_table(prefix="message", read_throughput=5, - write_throughput=5, delta=0): + write_throughput=5, delta=0, date=None): # type: (str, int, int, int) -> Table """Create a new message table for webpush style message storage""" - tablename = make_rotating_tablename(prefix, delta) + tablename = make_rotating_tablename(prefix, delta, date) return Table.create(tablename, schema=[HashKey("uaid"), RangeKey("chidmessageid")], @@ -153,7 +153,9 @@ def get_rotating_message_table(prefix="message", delta=0, date=None, return create_rotating_message_table( prefix=prefix, delta=delta, read_throughput=message_read_throughput, - write_throughput=message_write_throughput) + write_throughput=message_write_throughput, + date=date, + ) else: return Table(tablename) diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index 2f7ab146..a8f0f389 100644 --- a/autopush/tests/test_db.py +++ b/autopush/tests/test_db.py @@ -1,6 +1,6 @@ import unittest import uuid -from datetime import datetime +from datetime import datetime, timedelta from autopush.websocket import ms_time from boto.dynamodb2.exceptions import ( @@ -25,7 +25,7 @@ Message, Router, generate_last_connect, -) + make_rotating_tablename) from autopush.exceptions import AutopushException from autopush.metrics import SinkMetrics from autopush.utils import WebPushNotification @@ -360,6 +360,14 @@ def raise_condition(*args, **kwargs): result = message.delete_message(notif) eq_(result, False) + def test_message_rotate_table_with_date(self): + prefix = "message" + uuid.uuid4().hex + future = datetime.today() + timedelta(days=32) + tbl_name = make_rotating_tablename(prefix, date=future) + + m = get_rotating_message_table(prefix=prefix, date=future) + eq_(m.table_name, tbl_name) + class RouterTestCase(unittest.TestCase): def setUp(self):