icsahn
diff --git a/‎appengine/taskqueue/pull-counter/main.py
Copy file name to clipboardExpand all lines: appengine/taskqueue/pull-counter/main.py
+12-7Lines changed: 12 additions & 7 deletions b/‎appengine/taskqueue/pull-counter/main.py
Copy file name to clipboardExpand all lines: appengine/taskqueue/pull-counter/main.py
+12-7Lines changed: 12 additions & 7 deletions
diff --git a/‎appengine/taskqueue/pull-counter/pullcounter_test.py
Copy file name to clipboardExpand all lines: appengine/taskqueue/pull-counter/pullcounter_test.py
+9Lines changed: 9 additions & 0 deletions b/‎appengine/taskqueue/pull-counter/pullcounter_test.py
Copy file name to clipboardExpand all lines: appengine/taskqueue/pull-counter/pullcounter_test.py
+9Lines changed: 9 additions & 0 deletions
@@ -48,6 +48,13 @@ def post(self):
         self.redirect('/')
 
 
+@ndb.transactional
+def update_counter(key, tasks):
+    counter = Counter.get_or_insert(key, count=0)
+    counter.count += len(tasks)
+    counter.put()
+
+
 class CounterWorker(webapp2.RequestHandler):
     def get(self):
         """Indefinitely fetch tasks and update the datastore."""
@@ -60,20 +67,18 @@ def get(self):
                 logging.exception(e)
                 time.sleep(1)
                 continue
+
             if tasks:
                 key = tasks[0].tag
 
-                @ndb.transactional
-                def update_counter():
-                    counter = Counter.get_or_insert(key, count=0)
-                    counter.count += len(tasks)
-                    counter.put()
                 try:
-                    update_counter()
+                    update_counter(key, tasks)
                 except Exception as e:
                     logging.exception(e)
-                else:
+                    raise
+                finally:
                     queue.delete_tasks(tasks)
+
             time.sleep(1)
 
 
 
@@ -16,6 +16,7 @@
 
 from google.appengine.ext import testbed as gaetestbed
 import main
+import mock
 import webtest
 
 
@@ -31,3 +32,11 @@ def test_app(testbed):
     tasks = tq_stub.get_filtered_tasks()
     assert len(tasks) == 1
     assert tasks[0].name == 'task1'
+
+    with mock.patch('main.update_counter') as mock_update:
+        # Force update to fail, otherwise the loop will go forever.
+        mock_update.side_effect = RuntimeError()
+
+        app.get('/_ah/start', status=500)
+
+        assert mock_update.called