From d82d77430448794900fc026b49ff3182a803b866 Mon Sep 17 00:00:00 2001
From: Jon Wayne Parrott <jon.wayne.parrott@gmail.com>
Date: Mon, 16 May 2016 12:33:51 -0700
Subject: [PATCH] Adding test for taskqueue pull worker. Fixes #254

Change-Id: I65d943bd0ecba1251d46e1d028722bbf685e0251
---
 appengine/taskqueue/pull-counter/main.py      | 19 ++++++++++++-------
 .../pull-counter/pullcounter_test.py          |  9 +++++++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/appengine/taskqueue/pull-counter/main.py b/appengine/taskqueue/pull-counter/main.py
index 05d456cbf39a..0b021dc845ea 100644
--- a/appengine/taskqueue/pull-counter/main.py
+++ b/appengine/taskqueue/pull-counter/main.py
@@ -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)
 
 
diff --git a/appengine/taskqueue/pull-counter/pullcounter_test.py b/appengine/taskqueue/pull-counter/pullcounter_test.py
index 0614a8d3316e..15991cf6404a 100644
--- a/appengine/taskqueue/pull-counter/pullcounter_test.py
+++ b/appengine/taskqueue/pull-counter/pullcounter_test.py
@@ -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