Adding test for taskqueue pull worker. Fixes GoogleCloudPlatform#254 (GoogleCloudPlatform#349)

Author: Jon Wayne Parrott

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)
 
 

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