def test_update_rotating_tables(self):
from autopush.db import get_month
conf = AutopushConfig(hostname="example.com", resolve_hostname=True)
db = DatabaseManager.from_config(conf,
resource=autopush.tests.boto_resource)
db.create_initial_message_tables()
# Erase the tables it has on init, and move current month back one
last_month = get_month(-1)
db.current_month = last_month.month
db.message_tables = [
make_rotating_tablename("message", delta=-1),
make_rotating_tablename("message", delta=0)
]
# Create the next month's table, just in case today is the day before
# a new month, in which case the lack of keys will cause an error in
# update_rotating_tables
next_month = get_month(1)
assert next_month.month not in db.message_tables
# Get the deferred back
e = Deferred()
d = db.update_rotating_tables()
def check_tables(result):
assert db.current_month == get_month().month
assert len(db.message_tables) == 2
d.addCallback(check_tables)
d.addBoth(lambda x: e.callback(True))
return e
def test_migrate_user(self):
migrate_command = self._makeFUT()
# Create a user
last_month = make_rotating_tablename("message", delta=-1)
user = UserItemFactory(current_month=last_month)
uaid = user["uaid"]
self.db.router.register_user(user)
# Store some messages so we have some channels
self._store_messages(UUID(uaid), num=3)
# Check that it's there
item = self.db.router.get_uaid(uaid)
_, channels = self.db.message_tables[last_month].all_channels(uaid)
assert item["current_month"] != self.db.current_msg_month
assert item is not None
assert len(channels) == 3
# Migrate it
migrate_command.process(
MigrateUser(uaid=uaid, message_month=last_month))
# Check that it's in the new spot
item = self.db.router.get_uaid(uaid)
_, channels = self.db.message.all_channels(uaid)
assert item["current_month"] == self.db.current_msg_month
assert item is not None
assert len(channels) == 3
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)
def update_rotating_tables(self):
"""This method is intended to be tasked to run periodically off the
twisted event hub to rotate tables.
When today is a new month from yesterday, then we swap out all the
table objects on the settings object.
"""
today = datetime.date.today()
if today.month == self.current_month:
# No change in month, we're fine.
returnValue(False)
# Get tables for the new month, and verify they exist before we try to
# switch over
message_table = get_rotating_message_table(self._message_prefix)
try:
yield deferToThread(message_table.describe)
except Exception:
tblname = make_rotating_tablename(self._message_prefix)
log.err("Unable to locate new message table: %s" % tblname)
returnValue(False)
# Both tables found, safe to switch-over
self.current_month = today.month
self.current_msg_month = message_table.table_name
self.message_tables[self.current_msg_month] = \
Message(message_table, self.metrics)
returnValue(True)
def test_message_rotate_table_with_date(self):
prefix = "message" + uuid.uuid4().hex
future = (datetime.today() + timedelta(days=32)).date()
tbl_name = make_rotating_tablename(prefix, date=future)
m = get_rotating_message_table(prefix=prefix, date=future)
assert m.table_name == tbl_name
# Clean up the temp table.
_drop_table(tbl_name)
class StoreMessageFactory(factory.Factory):
class Meta:
model = StoreMessages
messages = factory.LazyAttribute(webpush_messages)
message_month = factory.LazyFunction(
lambda: make_rotating_tablename("message"))
class Params:
message_count = 20
uaid = factory.LazyFunction(lambda: uuid4().hex)
def test_webpush_monthly_rotation_no_channels(self):
from autopush.db import make_rotating_tablename
client = Client("ws://localhost:9010/", use_webpush=True)
yield client.connect()
yield client.hello()
yield client.disconnect()
# Move the client back one month to the past
last_month = make_rotating_tablename(
prefix=self._settings._message_prefix, delta=-1)
yield deferToThread(
self._settings.router.update_message_month,
client.uaid,
last_month
)
# Verify the move
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], last_month)
# Verify there's no channels
exists, chans = yield deferToThread(
self._settings.message.all_channels,
client.uaid
)
eq_(exists, False)
eq_(len(chans), 0)
# Connect the client, verify the migration
yield client.connect()
yield client.hello()
# Check that the client is going to rotate the month
server_client = self._settings.clients[client.uaid]
eq_(server_client.ps.rotate_message_table, True)
# Wait up to 2 seconds for the table rotation to occur
start = time.time()
while time.time()-start < 2:
c = yield deferToThread(self._settings.router.get_uaid,
client.uaid)
if c["current_month"] == self._settings.current_msg_month:
break
else:
yield deferToThread(time.sleep, 0.2)
# Verify the month update in the router table
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], self._settings.current_msg_month)
eq_(server_client.ps.rotate_message_table, False)
yield self.shut_down(client)
def test_webpush_monthly_rotation_no_channels(self):
from autopush.db import make_rotating_tablename
client = Client("ws://localhost:9010/", use_webpush=True)
yield client.connect()
yield client.hello()
yield client.disconnect()
# Move the client back one month to the past
last_month = make_rotating_tablename(
prefix=self._settings._message_prefix, delta=-1)
yield deferToThread(
self._settings.router.update_message_month,
client.uaid,
last_month
)
# Verify the move
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], last_month)
# Verify there's no channels
exists, chans = yield deferToThread(
self._settings.message.all_channels,
client.uaid
)
eq_(exists, False)
eq_(len(chans), 0)
# Connect the client, verify the migration
yield client.connect()
yield client.hello()
# Check that the client is going to rotate the month
server_client = self._settings.clients[client.uaid]
eq_(server_client.ps.rotate_message_table, True)
# Wait up to 2 seconds for the table rotation to occur
start = time.time()
while time.time()-start < 2:
c = yield deferToThread(self._settings.router.get_uaid,
client.uaid)
if c["current_month"] == self._settings.current_msg_month:
break
else:
yield deferToThread(time.sleep, 0.2)
# Verify the month update in the router table
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], self._settings.current_msg_month)
eq_(server_client.ps.rotate_message_table, False)
yield self.shut_down(client)
class UserItemFactory(factory.Factory):
class Meta:
model = dict
uaid = factory.LazyFunction(lambda: uuid4().hex)
connected_at = factory.LazyFunction(
lambda: int(time.time() * 1000) - 10000)
node_id = "http://something:3242/"
router_type = "webpush"
last_connect = factory.LazyFunction(generate_last_connect)
record_version = USER_RECORD_VERSION
current_month = factory.LazyFunction(
lambda: make_rotating_tablename("message"))
def test_webpush_monthly_rotation(self):
from autopush.db import make_rotating_tablename
client = yield self.quick_register(use_webpush=True)
yield client.disconnect()
# Move the client back one month to the past
last_month = make_rotating_tablename(prefix=self._settings._message_prefix, delta=-1)
lm_message = self._settings.message_tables[last_month]
yield deferToThread(self._settings.router.update_message_month, client.uaid, last_month)
# Verify the move
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], last_month)
# Verify last_connect is current, then move that back
ok_(has_connected_this_month(c))
today = get_month(delta=-1)
c["last_connect"] = "%s%s020001" % (today.year, str(today.month).zfill(2))
yield deferToThread(c.partial_save)
eq_(False, has_connected_this_month(c))
# Move the clients channels back one month
exists, chans = yield deferToThread(self._settings.message.all_channels, client.uaid)
eq_(exists, True)
eq_(len(chans), 1)
yield deferToThread(lm_message.save_channels, client.uaid, chans)
# Remove the channels entry entirely from this month
yield deferToThread(self._settings.message.table.delete_item, uaid=client.uaid, chidmessageid=" ")
# Verify the channel is gone
exists, chans = yield deferToThread(self._settings.message.all_channels, client.uaid)
eq_(exists, False)
eq_(len(chans), 0)
# Send in a notification, verify it landed in last months notification
# table
data = uuid.uuid4().hex
yield client.send_notification(data=data)
notifs = yield deferToThread(lm_message.fetch_messages, client.uaid)
eq_(len(notifs), 1)
# Connect the client, verify the migration
yield client.connect()
yield client.hello()
# Pull down the notification
result = yield client.get_notification()
chan = client.channels.keys()[0]
ok_(result is not None)
eq_(chan, result["channelID"])
# Check that the client is going to rotate the month
server_client = self._settings.clients[client.uaid]
eq_(server_client.ps.rotate_message_table, True)
# Acknowledge the notification, which triggers the migration
yield client.ack(chan, result["version"])
# Wait up to 2 seconds for the table rotation to occur
start = time.time()
while time.time() - start < 2:
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
if c["current_month"] == self._settings.current_msg_month:
break
else:
yield deferToThread(time.sleep, 0.2)
# Verify the month update in the router table
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], self._settings.current_msg_month)
eq_(server_client.ps.rotate_message_table, False)
# Verify the client moved last_connect
eq_(True, has_connected_this_month(c))
# Verify the channels were moved
exists, chans = yield deferToThread(self._settings.message.all_channels, client.uaid)
eq_(exists, True)
eq_(len(chans), 1)
yield self.shut_down(client)
def test_zero_fill_month(self):
from autopush.db import make_rotating_tablename
assert 'test_2016_03' == make_rotating_tablename('test',
date=datetime(
2016, 3,
15).date())
def __init__(self,
crypto_key=None,
datadog_api_key=None,
datadog_app_key=None,
datadog_flush_interval=None,
hostname=None,
port=None,
router_scheme=None,
router_hostname=None,
router_port=None,
endpoint_scheme=None,
endpoint_hostname=None,
endpoint_port=None,
router_conf={},
router_tablename="router",
router_read_throughput=5,
router_write_throughput=5,
storage_tablename="storage",
storage_read_throughput=5,
storage_write_throughput=5,
message_tablename="message",
message_read_throughput=5,
message_write_throughput=5,
statsd_host="localhost",
statsd_port=8125,
resolve_hostname=False,
max_data=4096,
# Reflected up from UDP Router
wake_timeout=0,
env='development',
enable_cors=False,
s3_bucket=DEFAULT_BUCKET,
senderid_expry=SENDERID_EXPRY,
senderid_list={},
hello_timeout=0,
auth_key=None,
):
"""Initialize the Settings object
Upon creation, the HTTP agent will initialize, all configured routers
will be setup and started, logging will be started, and the database
will have a preflight check done.
"""
# Use a persistent connection pool for HTTP requests.
pool = HTTPConnectionPool(reactor)
self.agent = Agent(reactor, connectTimeout=5, pool=pool)
# Metrics setup
if datadog_api_key:
self.metrics = DatadogMetrics(
api_key=datadog_api_key,
app_key=datadog_app_key,
flush_interval=datadog_flush_interval
)
elif statsd_host:
self.metrics = TwistedMetrics(statsd_host, statsd_port)
else:
self.metrics = SinkMetrics()
if not crypto_key:
crypto_key = [Fernet.generate_key()]
if not isinstance(crypto_key, list):
crypto_key = [crypto_key]
self.update(crypto_key=crypto_key)
self.crypto_key = crypto_key
if auth_key is None:
auth_key = []
if not isinstance(auth_key, list):
auth_key = [auth_key]
self.auth_key = auth_key
self.max_data = max_data
self.clients = {}
# Setup hosts/ports/urls
default_hostname = socket.gethostname()
self.hostname = hostname or default_hostname
if resolve_hostname:
self.hostname = resolve_ip(self.hostname)
self.port = port
self.endpoint_hostname = endpoint_hostname or self.hostname
self.router_hostname = router_hostname or self.hostname
self.router_conf = router_conf
self.router_url = canonical_url(
router_scheme or 'http',
self.router_hostname,
router_port
)
self.endpoint_url = canonical_url(
endpoint_scheme or 'http',
self.endpoint_hostname,
endpoint_port
)
# Database objects
self.router_table = get_router_table(router_tablename,
router_read_throughput,
router_write_throughput)
self.storage_table = get_storage_table(
storage_tablename,
storage_read_throughput,
storage_write_throughput)
self.message_table = get_rotating_message_table(
message_tablename)
self._message_prefix = message_tablename
self.storage = Storage(self.storage_table, self.metrics)
self.router = Router(self.router_table, self.metrics)
# Used to determine whether a connection is out of date with current
# db objects
self.current_msg_month = make_rotating_tablename(self._message_prefix)
self.current_month = datetime.date.today().month
self.create_initial_message_tables()
# Run preflight check
preflight_check(self.storage, self.router)
# CORS
self.cors = enable_cors
# Force timeout in idle seconds
self.wake_timeout = wake_timeout
# Setup the routers
self.routers = {}
self.routers["simplepush"] = SimpleRouter(
self,
router_conf.get("simplepush")
)
self.routers["webpush"] = WebPushRouter(self, None)
if 'apns' in router_conf:
self.routers["apns"] = APNSRouter(self, router_conf["apns"])
if 'gcm' in router_conf:
self.routers["gcm"] = GCMRouter(self, router_conf["gcm"])
# Env
self.env = env
self.hello_timeout = hello_timeout
def test_zero_fill_month(self):
from autopush.db import make_rotating_tablename
eq_('test_2016_03',
make_rotating_tablename('test', date=datetime(2016, 3, 15)))
def test_webpush_monthly_rotation_prior_record_exists(self):
from autopush.db import make_rotating_tablename
client = yield self.quick_register(use_webpush=True)
yield client.disconnect()
# Move the client back one month to the past
last_month = make_rotating_tablename(
prefix=self._settings._message_prefix, delta=-1)
lm_message = self._settings.message_tables[last_month]
yield deferToThread(
self._settings.router.update_message_month,
client.uaid,
last_month
)
# Verify the move
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], last_month)
# Verify last_connect is current, then move that back
ok_(has_connected_this_month(c))
today = get_month(delta=-1)
c["last_connect"] = "%s%s020001" % (today.year,
str(today.month).zfill(2))
yield deferToThread(c.partial_save)
eq_(False, has_connected_this_month(c))
# Move the clients channels back one month
exists, chans = yield deferToThread(
self._settings.message.all_channels, client.uaid
)
eq_(exists, True)
eq_(len(chans), 1)
yield deferToThread(
lm_message.save_channels,
client.uaid,
chans
)
# Send in a notification, verify it landed in last months notification
# table
data = uuid.uuid4().hex
yield client.send_notification(data=data)
notifs = yield deferToThread(lm_message.fetch_messages, client.uaid)
eq_(len(notifs), 1)
# Connect the client, verify the migration
yield client.connect()
yield client.hello()
# Pull down the notification
result = yield client.get_notification()
chan = client.channels.keys()[0]
ok_(result is not None)
eq_(chan, result["channelID"])
# Check that the client is going to rotate the month
server_client = self._settings.clients[client.uaid]
eq_(server_client.ps.rotate_message_table, True)
# Acknowledge the notification, which triggers the migration
yield client.ack(chan, result["version"])
# Wait up to 2 seconds for the table rotation to occur
start = time.time()
while time.time()-start < 2:
c = yield deferToThread(self._settings.router.get_uaid,
client.uaid)
if c["current_month"] == self._settings.current_msg_month:
break
else:
yield deferToThread(time.sleep, 0.2)
# Verify the month update in the router table
c = yield deferToThread(self._settings.router.get_uaid, client.uaid)
eq_(c["current_month"], self._settings.current_msg_month)
eq_(server_client.ps.rotate_message_table, False)
# Verify the client moved last_connect
eq_(True, has_connected_this_month(c))
# Verify the channels were moved
exists, chans = yield deferToThread(
self._settings.message.all_channels,
client.uaid
)
eq_(exists, True)
eq_(len(chans), 1)
yield self.shut_down(client)
