diff --git a/autopush/db.py b/autopush/db.py index a34781f1..9b48022c 100644 --- a/autopush/db.py +++ b/autopush/db.py @@ -18,6 +18,7 @@ from boto.dynamodb2.table import Table from boto.dynamodb2.types import NUMBER +from autopush.exceptions import AutopushException from autopush.utils import generate_hash key_hash = "" @@ -582,6 +583,10 @@ def register_user(self, data): conn = self.table.connection db_key = self.encode({"uaid": hasher(data["uaid"])}) del data["uaid"] + if "router_type" not in data or "connected_at" not in data: + # Not specifying these values will generate an exception in AWS. + raise AutopushException("data is missing router_type" + "or connected_at") # Generate our update expression expr = "SET " + ", ".join(["%s=:%s" % (x, x) for x in data.keys()]) expr_values = self.encode({":%s" % k: v for k, v in data.items()}) diff --git a/autopush/endpoint.py b/autopush/endpoint.py index 38c4404b..89e38065 100644 --- a/autopush/endpoint.py +++ b/autopush/endpoint.py @@ -569,12 +569,8 @@ def _router_completed(self, response, uaid_data, warning=""): # TODO: Add some custom wake logic here # Were we told to update the router data? - if response.router_data is not None: - if not response.router_data: - del uaid_data["router_data"] - del uaid_data["router_type"] - else: - uaid_data["router_data"] = response.router_data + if response.router_data: + uaid_data["router_data"] = response.router_data uaid_data["connected_at"] = int(time.time() * 1000) d = deferToThread(self.ap_settings.router.register_user, uaid_data) diff --git a/autopush/router/gcm.py b/autopush/router/gcm.py index e17b1716..223aaf22 100644 --- a/autopush/router/gcm.py +++ b/autopush/router/gcm.py @@ -94,9 +94,9 @@ def _route(self, notification, router_data): except KeyError: raise self._error("Server error, missing bridge credentials " + "for %s" % creds.get("senderID"), 500) - except gcmclient.GCMAuthenticationError, e: + except gcmclient.GCMAuthenticationError as e: raise self._error("Authentication Error: %s" % e, 500) - except Exception, e: + except Exception as e: raise self._error("Unhandled exception in GCM Routing: %s" % e, 500) self.metrics.increment("updates.client.bridge.gcm.attempted", diff --git a/autopush/tests/test_db.py b/autopush/tests/test_db.py index d019cf2c..3494955a 100644 --- a/autopush/tests/test_db.py +++ b/autopush/tests/test_db.py @@ -10,7 +10,7 @@ from boto.dynamodb2.layer1 import DynamoDBConnection from boto.dynamodb2.items import Item from mock import Mock -from nose.tools import eq_, assert_raises +from nose.tools import eq_, assert_raises, ok_ from autopush.db import ( get_rotating_message_table, @@ -459,7 +459,8 @@ def raise_error(*args, **kwargs): router.table.connection.update_item.side_effect = raise_error with self.assertRaises(ProvisionedThroughputExceededException): router.register_user(dict(uaid=dummy_uaid, node_id="me", - connected_at=1234)) + connected_at=1234, + router_type="simplepush")) def test_clear_node_provision_failed(self): r = get_router_table() @@ -473,19 +474,24 @@ def raise_error(*args, **kwargs): with self.assertRaises(ProvisionedThroughputExceededException): router.clear_node(Item(r, dict(uaid=dummy_uaid, connected_at="1234", - node_id="asdf"))) + node_id="asdf", + router_type="simplepush"))) def test_incomplete_uaid(self): + # Older records may be incomplete. We can't inject them using normal + # methods. uaid = str(uuid.uuid4()) r = get_router_table() router = Router(r, SinkMetrics()) + router.table.get_item = Mock() + router.drop_user = Mock() + router.table.get_item.return_value = {"uaid": uuid.uuid4().hex} try: router.register_user(dict(uaid=uaid)) except: pass self.assertRaises(ItemNotFound, router.get_uaid, uaid) - self.assertRaises(ItemNotFound, router.table.get_item, - consistent=True, uaid=uaid) + ok_(router.drop_user.called) def test_save_new(self): r = get_router_table() @@ -495,6 +501,7 @@ def test_save_new(self): router.table.connection = Mock() router.table.connection.update_item.return_value = {} result = router.register_user(dict(uaid="", node_id="me", + router_type="simplepush", connected_at=1234)) eq_(result[0], True) @@ -507,7 +514,8 @@ def raise_condition(*args, **kwargs): router.table.connection = Mock() router.table.connection.update_item.side_effect = raise_condition - router_data = dict(uaid=dummy_uaid, node_id="asdf", connected_at=1234) + router_data = dict(uaid=dummy_uaid, node_id="asdf", connected_at=1234, + router_type="simplepush") result = router.register_user(router_data) eq_(result, (False, {}, router_data)) @@ -519,7 +527,6 @@ def test_node_clear(self): router.register_user(dict(uaid=dummy_uaid, node_id="asdf", connected_at=1234, router_type="webpush")) - # Verify user = router.get_uaid(dummy_uaid) eq_(user["node_id"], "asdf") @@ -553,8 +560,9 @@ def test_drop_user(self): r = get_router_table() router = Router(r, SinkMetrics()) # Register a node user - router.table.put_item(data=dict(uaid=uaid, node_id="asdf", - connected_at=1234)) + router.register_user(dict(uaid=uaid, node_id="asdf", + router_type="simplepush", + connected_at=1234)) result = router.drop_user(uaid) eq_(result, True) # Deleting already deleted record should return false. diff --git a/autopush/tests/test_endpoint.py b/autopush/tests/test_endpoint.py index 292900a3..65a07171 100644 --- a/autopush/tests/test_endpoint.py +++ b/autopush/tests/test_endpoint.py @@ -624,7 +624,7 @@ def test_put_router_needs_change(self): def handle_finish(result): self.assertTrue(result) self.endpoint.set_status.assert_called_with(500, None) - assert(self.router_mock.register_user.called) + ok_(not self.router_mock.register_user.called) self.finish_deferred.addCallback(handle_finish) self.endpoint.put(None, dummy_uaid) diff --git a/autopush/tests/test_websocket.py b/autopush/tests/test_websocket.py index bb0f4ab3..c9a7d24c 100644 --- a/autopush/tests/test_websocket.py +++ b/autopush/tests/test_websocket.py @@ -435,6 +435,7 @@ def test_hello_old(self): uaid=orig_uaid, connected_at=ms_time(), current_month=msg_date, + router_type="webpush" )) def fake_msg(data): @@ -1866,6 +1867,19 @@ def after_hello(result): self.proto.ps._register.addErrback(lambda x: d.errback(x)) return d + def test_incomplete_uaid(self): + mm = self.proto.ap_settings.router = Mock() + fr = self.proto.force_retry = Mock() + uaid = uuid.uuid4().hex + mm.get_uaid.return_value = { + 'uaid': uaid + } + self.proto.ps.uaid = uaid + reply = self.proto._verify_user_record() + eq_(reply, None) + assert(fr.called) + eq_(fr.call_args[0], (mm.drop_user, uaid)) + class RouterHandlerTestCase(unittest.TestCase): def setUp(self):