swarming: mechanical changes to unit tests

appengine/swarming/server/task_scheduler_test.py

 #!/usr/bin/env python
 # Copyright 2014 The LUCI Authors. All rights reserved.
 # Use of this source code is governed under the Apache License, Version 2.0
 # that can be found in the LICENSE file.
 
 import datetime
 import logging
 import os
 import random
 import sys
@@ skipping 29 matching lines @@
 from server import task_scheduler
 from server import task_to_run
 from server.task_result import State
 
 from proto import config_pb2
 
 
 # pylint: disable=W0212,W0612
 
 
-def _gen_request(properties=None, **kwargs):
-  """Creates a TaskRequest."""
-  props = {
-    'command': [u'command1'],
-    'dimensions': {u'pool': u'default'},
-    'env': {},
-    'execution_timeout_secs': 24*60*60,
-    'io_timeout_secs': None,
-  }
-  props.update(properties or {})
-  now = utils.utcnow()
-  args = {
-    'created_ts': now,
-    'name': 'Request name',
-    'priority': 50,
-    'properties': task_request.TaskProperties(**props),
-    'expiration_ts': now + datetime.timedelta(seconds=60),
-    'tags': [u'tag:1'],
-    'user': 'Jesus',
-  }
-  args.update(kwargs)
-  ret = task_request.TaskRequest(**args)
-  task_request.init_new_request(ret, True, None)
-  return ret
-
-
 def get_results(request_key):
   """Fetches all task results for a specified TaskRequest ndb.Key.
 
   Returns:
     tuple(TaskResultSummary, list of TaskRunResult that exist).
   """
   result_summary_key = task_pack.request_key_to_result_summary_key(request_key)
   result_summary = result_summary_key.get()
   # There's two way to look at it, either use a DB query or fetch all the
   # entities that could exist, at most 255. In general, there will be <3
   # entities so just fetching them by key would be faster. This function is
   # exclusively used in unit tests so it's not performance critical.
   q = task_result.TaskRunResult.query(ancestor=result_summary_key)
   q = q.order(task_result.TaskRunResult.key)
   return result_summary, q.fetch()
 
 
-def _quick_schedule(dims):
-  """Schedules a task."""
-  request = _gen_request(properties={'dimensions': dims})
-  task_request.init_new_request(request, True, None)
-  return task_scheduler.schedule_request(request, None)
-
-
-def _register_bot(bot_dimensions):
-  """Registers the bot so the task queues knows there's a worker than can run
-  the task.
-  """
-  bot_management.bot_event(
-      'bot_connected', bot_dimensions[u'id'][0], '1.2.3.4', 'joe@localhost',
-      bot_dimensions, {'state': 'real'}, '1234', False, None, None)
-  task_queues.assert_bot(bot_dimensions)
-
-
-def _quick_reap():
-  """Reaps a task."""
-  _quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'})
-  bot_dimensions = {
-    u'id': [u'localhost'],
-    u'os': [u'Windows-3.1.1'],
-    u'pool': [u'default'],
-  }
-  _register_bot(bot_dimensions)
-  reaped_request, _, run_result = task_scheduler.bot_reap_task(
-      bot_dimensions, 'abc', None)
-  return run_result
-
-
 class TaskSchedulerApiTest(test_env_handlers.AppTestBase):
   def setUp(self):
     super(TaskSchedulerApiTest, self).setUp()
     self.now = datetime.datetime(2014, 1, 2, 3, 4, 5, 6)
     self.mock_now(self.now)
     self.mock(stats_framework, 'add_entry', self._parse_line)
     auth_testing.mock_get_current_identity(self)
     event_mon_metrics.initialize()
     # Setup the backend to handle task queues for 'task-dimensions'.
     self.app = webtest.TestApp(
@@ skipping 26 matching lines @@
     self.assertFalse(self._pub_sub_mocked)
     self._pub_sub_mocked = True
     calls = []
     def pubsub_publish(**kwargs):
       if not self.publish_successful:
         raise pubsub.TransientError('Fail')
       calls.append(('directly', kwargs))
     self.mock(pubsub, 'publish', pubsub_publish)
     return calls
 
+  def _gen_request(self, properties=None, **kwargs):
+    """Creates a TaskRequest."""
+    props = {
+      'command': [u'command1'],
+      'dimensions': {u'pool': u'default'},
+      'env': {},
+      'execution_timeout_secs': 24*60*60,
+      'io_timeout_secs': None,
+    }
+    props.update(properties or {})
+    now = utils.utcnow()
+    args = {
+      'created_ts': now,
+      'name': 'Request name',
+      'priority': 50,
+      'properties': task_request.TaskProperties(**props),
+      'expiration_ts': now + datetime.timedelta(seconds=60),
+      'tags': [u'tag:1'],
+      'user': 'Jesus',
+    }
+    args.update(kwargs)
+    ret = task_request.TaskRequest(**args)
+    task_request.init_new_request(ret, True, None)
+    return ret
+
+  def _quick_schedule(self, dims, nb_task=0):
+    """Schedules a task."""
+    request = self._gen_request(properties={'dimensions': dims})
+    result_summary = task_scheduler.schedule_request(request, None)
+    self.assertEqual(nb_task, self.execute_tasks())
+    return result_summary
+
+  def _register_bot(self, bot_dimensions, nb_task=0):
+    """Registers the bot so the task queues knows there's a worker than can run
+    the task.
+    """
+    bot_management.bot_event(
+        'bot_connected', bot_dimensions[u'id'][0], '1.2.3.4', 'joe@localhost',
+        bot_dimensions, {'state': 'real'}, '1234', False, None, None)
+    task_queues.assert_bot(bot_dimensions)
+    self.assertEqual(nb_task, self.execute_tasks())
+
+  def _quick_reap(self, nb_task=0):
+    """Reaps a task."""
+    self._quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'})
+    bot_dimensions = {
+      u'id': [u'localhost'],
+      u'os': [u'Windows-3.1.1'],
+      u'pool': [u'default'],
+    }
+    self._register_bot(bot_dimensions, nb_task=nb_task)
+    reaped_request, _, run_result = task_scheduler.bot_reap_task(
+        bot_dimensions, 'abc', None)
+    return run_result
+
   def test_all_apis_are_tested(self):
     # Ensures there's a test for each public API.
     # TODO(maruel): Remove this once coverage is asserted.
     module = task_scheduler
     expected = set(
         i for i in dir(module)
         if i[0] != '_' and hasattr(getattr(module, i), 'func_name'))
     missing = expected - set(i[5:] for i in dir(self) if i.startswith('test_'))
     self.assertFalse(missing)
 
   def test_bot_reap_task(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     actual_request, _, run_result  = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(request, actual_request)
     self.assertEqual('localhost', run_result.bot_id)
     self.assertEqual(None, task_to_run.TaskToRun.query().get().queue_number)
 
   def test_bot_reap_task_not_enough_time(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     actual_request, _, run_result  = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', datetime.datetime(1969, 1, 1))
     self.failIf(actual_request)
     self.failIf(run_result)
     self.failUnless(task_to_run.TaskToRun.query().get().queue_number)
 
   def test_bot_reap_task_enough_time(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     actual_request, _, run_result  = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', datetime.datetime(3000, 1, 1))
     self.assertEqual(request, actual_request)
     self.assertEqual('localhost', run_result.bot_id)
     self.failIf(task_to_run.TaskToRun.query().get().queue_number)
 
   def test_exponential_backoff(self):
     self.mock(
         task_scheduler.random, 'random',
         lambda: task_scheduler._PROBABILITY_OF_QUICK_COMEBACK)
@@ skipping 35 matching lines @@
     })
     self.assertEqual([
       {
         'attributes': {'auth_token': 'token'},
         'message': '{"task_id":"abcdef123","userdata":"userdata"}',
         'topic': 'projects/abc/topics/def',
     }], calls)
 
   def _task_ran_successfully(self):
     """Runs a task successfully and returns the task_id."""
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     actual_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(request, actual_request)
     self.assertEqual('localhost', run_result.bot_id)
     self.assertEqual(None, task_to_run.TaskToRun.query().get().queue_number)
     # It's important to terminate the task with success.
     self.assertEqual(
         task_result.State.COMPLETED,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='Foo1',
             output_chunk_start=0,
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
     return unicode(run_result.task_id)
 
-  def _task_deduped(self, new_ts, deduped_from, task_id, nb_queues, now=None):
-    request = _gen_request(
+  def _task_deduped(self, new_ts, deduped_from, task_id, nb_task=0, now=None):
+    request = self._gen_request(
         name='yay',
         user='Raoul',
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     self.assertEqual(None, task_to_run.TaskToRun.query().get().queue_number)
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions, nb_task=nb_task)
     actual_request_2, _, run_result_2 = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(None, actual_request_2)
-    self.assertEqual(nb_queues, self.execute_tasks())
     result_summary_duped, run_results_duped = get_results(request.key)
     expected = {
       'abandoned_ts': None,
       'bot_dimensions': bot_dimensions,
       'bot_id': u'localhost',
       'bot_version': u'abc',
       'cipd_pins': None,
       'children_task_ids': [],
       'completed_ts': now or self.now,
       'costs_usd': [],
@@ skipping 28 matching lines @@
     }
     self.assertEqual(expected, result_summary_duped.to_dict())
     self.assertEqual([], run_results_duped)
 
   def test_task_idempotent(self):
     # First task is idempotent.
     task_id = self._task_ran_successfully()
 
     # Second task is deduped against first task.
     new_ts = self.mock_now(self.now, config.settings().reusable_task_age_secs-1)
-    self._task_deduped(new_ts, task_id, '1d8dc670a0008a10', 0)
+    self._task_deduped(new_ts, task_id, '1d8dc670a0008a10')
 
   def test_task_idempotent_old(self):
     # First task is idempotent.
     self._task_ran_successfully()
 
     # Second task is scheduled, first task is too old to be reused.
     new_ts = self.mock_now(self.now, config.settings().reusable_task_age_secs)
-    request = _gen_request(
+    request = self._gen_request(
         name='yay',
         user='Raoul',
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     # The task was enqueued for execution.
     self.assertNotEqual(None, task_to_run.TaskToRun.query().get().queue_number)
 
   def test_task_idempotent_three(self):
     # First task is idempotent.
     task_id = self._task_ran_successfully()
 
     # Second task is deduped against first task.
     new_ts = self.mock_now(self.now, config.settings().reusable_task_age_secs-1)
-    self._task_deduped(new_ts, task_id, '1d8dc670a0008a10', 0)
+    self._task_deduped(new_ts, task_id, '1d8dc670a0008a10')
 
     # Third task is scheduled, second task is not dedupable, first task is too
     # old.
     new_ts = self.mock_now(self.now, config.settings().reusable_task_age_secs)
-    request = _gen_request(
+    request = self._gen_request(
         name='yay',
         user='Jesus',
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     # The task was enqueued for execution.
     self.assertNotEqual(None, task_to_run.TaskToRun.query().get().queue_number)
@@ skipping 12 matching lines @@
     second_ts = self.mock_now(self.now, 10)
     task_id = self._task_ran_successfully()
 
     # Now any of the 2 tasks could be reused. Assert the right one (the most
     # recent) is reused.
     cfg.reusable_task_age_secs = 100
 
     # Third task is deduped against second task. That ensures ordering works
     # correctly.
     third_ts = self.mock_now(self.now, 20)
-    self._task_deduped(third_ts, task_id, '1d69ba3ea8008b10', 0, second_ts)
+    self._task_deduped(third_ts, task_id, '1d69ba3ea8008b10', now=second_ts)
 
   def test_task_parent_children(self):
     # Parent task creates a child task.
     parent_id = self._task_ran_successfully()
-    request = _gen_request(
+    request = self._gen_request(
         parent_task_id=parent_id,
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     self.assertEqual([], result_summary.children_task_ids)
     self.assertEqual(parent_id, request.parent_task_id)
 
     parent_run_result_key = task_pack.unpack_run_result_key(parent_id)
     parent_res_summary_key = task_pack.run_result_key_to_result_summary_key(
         parent_run_result_key)
     expected = [result_summary.task_id]
     self.assertEqual(expected, parent_run_result_key.get().children_task_ids)
     self.assertEqual(expected, parent_res_summary_key.get().children_task_ids)
 
   def test_task_parent_isolated(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'command': [],
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'inputs_ref': {
             'isolated': '1' * 40,
             'isolatedserver': 'http://localhost:1',
             'namespace': 'default-gzip',
           },
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     actual_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(request, actual_request)
     self.assertEqual('localhost', run_result.bot_id)
     self.assertEqual(None, task_to_run.TaskToRun.query().get().queue_number)
     # It's important to terminate the task with success.
     self.assertEqual(
         task_result.State.COMPLETED,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='Foo1',
             output_chunk_start=0,
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
 
     parent_id = run_result.task_id
-    request = _gen_request(
+    request = self._gen_request(
         parent_task_id=parent_id,
         properties={
           'dimensions':{u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     self.assertEqual([], result_summary.children_task_ids)
     self.assertEqual(parent_id, request.parent_task_id)
 
     parent_run_result_key = task_pack.unpack_run_result_key(parent_id)
     parent_res_summary_key = task_pack.run_result_key_to_result_summary_key(
         parent_run_result_key)
     expected = [result_summary.task_id]
     self.assertEqual(expected, parent_run_result_key.get().children_task_ids)
     self.assertEqual(expected, parent_res_summary_key.get().children_task_ids)
 
   def test_get_results(self):
     # TODO(maruel): Split in more focused tests.
     created_ts = self.now
     self.mock_now(created_ts)
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
 
     # The TaskRequest was enqueued, the TaskResultSummary was created but no
     # TaskRunResult exist yet since the task was not scheduled on any bot.
     result_summary, run_results = get_results(request.key)
     expected = {
@@ skipping 35 matching lines @@
     self.assertEqual([], run_results)
 
     # A bot reaps the TaskToRun.
     reaped_ts = self.now + datetime.timedelta(seconds=60)
     self.mock_now(reaped_ts)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(request, reaped_request)
     self.assertTrue(run_result)
     result_summary, run_results = get_results(request.key)
     expected = {
       'abandoned_ts': None,
       'bot_dimensions': bot_dimensions,
       'bot_id': u'localhost',
       'bot_version': u'abc',
@@ skipping 166 matching lines @@
         },
         'server_versions': [u'v1a'],
         'started_ts': reaped_ts,
         'state': State.COMPLETED,
         'try_number': 1,
       },
     ]
     self.assertEqual(expected, [t.to_dict() for t in run_results])
 
   def test_exit_code_failure(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(request, reaped_request)
     self.assertEqual(
         task_result.State.COMPLETED,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='Foo1',
@@ skipping 65 matching lines @@
         'started_ts': self.now,
         'state': State.COMPLETED,
         'try_number': 1,
       },
     ]
     self.assertEqual(expected, [t.to_dict() for t in run_results])
 
   def test_schedule_request(self):
     # It is tested indirectly in the other functions.
     self.assertTrue(
-        _quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'}))
+        self._quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'}))
 
   def mock_dim_acls(self, mapping):
     self.mock(config, 'settings', lambda: config_pb2.SettingsCfg(
       dimension_acls=config_pb2.DimensionACLs(entry=[
         config_pb2.DimensionACLs.Entry(dimension=[d], usable_by=g)
         for d, g in sorted(mapping.iteritems())
       ]),
     ))
 
   def test_schedule_request_forbidden_dim(self):
     self.mock_dim_acls({u'pool:bad': u'noone'})
-    _quick_schedule({u'pool': u'good'})
+    self._quick_schedule({u'pool': u'good'})
     with self.assertRaises(auth.AuthorizationError):
-      _quick_schedule({u'pool': u'bad'})
+      self._quick_schedule({u'pool': u'bad'})
 
   def test_schedule_request_forbidden_dim_via_star(self):
     self.mock_dim_acls({u'abc:*': u'noone'})
-    _quick_schedule({u'pool': u'default'})
+    self._quick_schedule({u'pool': u'default'})
     with self.assertRaises(auth.AuthorizationError):
-      _quick_schedule({u'pool': u'default', u'abc': u'blah'})
+      self._quick_schedule({u'pool': u'default', u'abc': u'blah'})
 
   def test_schedule_request_id_without_pool(self):
     self.mock_dim_acls({u'pool:good': u'mocked'})
     with self.assertRaises(auth.AuthorizationError):
-      _quick_schedule({u'id': u'abc'})
+      self._quick_schedule({u'id': u'abc'})
     auth_testing.mock_is_admin(self)
-    _quick_schedule({u'id': u'abc'})
+    self._quick_schedule({u'id': u'abc'})
 
   def test_schedule_request_id_and_pool(self):
     self.mock_dim_acls({u'pool:good': u'mocked'})
     self.mock_dim_acls({u'pool:bad': u'unknown'})
 
     def mocked_is_group_member(group, ident):
       if group == 'mocked' and ident == auth_testing.DEFAULT_MOCKED_IDENTITY:
         return True
       return False
     self.mock(auth, 'is_group_member', mocked_is_group_member)
 
-    _quick_schedule({u'id': u'abc', u'pool': u'unknown'})
-    _quick_schedule({u'id': u'abc', u'pool': u'good'})
+    self._quick_schedule({u'id': u'abc', u'pool': u'unknown'})
+    self._quick_schedule({u'id': u'abc', u'pool': u'good'})
     with self.assertRaises(auth.AuthorizationError):
-      _quick_schedule({u'id': u'abc', u'pool': u'bad'})
+      self._quick_schedule({u'id': u'abc', u'pool': u'bad'})
 
   def test_bot_update_task(self):
-    run_result = _quick_reap()
+    run_result = self._quick_reap()
     self.assertEqual(
         task_result.State.RUNNING,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='hi',
             output_chunk_start=0,
             exit_code=None,
             duration=None,
@@ skipping 13 matching lines @@
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
     self.assertEqual('hihey', run_result.key.get().get_output())
 
   def test_bot_update_task_new_overwrite(self):
-    run_result = _quick_reap()
+    run_result = self._quick_reap()
     self.assertEqual(
         task_result.State.RUNNING,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='hi',
             output_chunk_start=0,
             exit_code=None,
             duration=None,
@@ skipping 13 matching lines @@
             exit_code=None,
             duration=None,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
     self.assertEqual('hhey', run_result.key.get().get_output())
 
   def test_bot_update_exception(self):
-    run_result = _quick_reap()
+    run_result = self._quick_reap()
     def r(*_):
       raise datastore_utils.CommitError('Sorry!')
 
     self.mock(ndb, 'put_multi', r)
     self.assertEqual(
         None,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='hi',
             output_chunk_start=0,
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
 
   def test_bot_update_pubsub_error(self):
     pub_sub_calls = self.mock_pub_sub()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         },
         pubsub_topic='projects/abc/topics/def')
     task_request.init_new_request(request, True, None)
     task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     _, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual('localhost', run_result.bot_id)
-    self.assertEqual(1, self.execute_tasks())
 
     # Attempt to terminate the task with success, but make PubSub call fail.
     self.publish_successful = False
     self.assertEqual(
         None,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='Foo1',
             output_chunk_start=0,
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
-    self.assertEqual(0, self.execute_tasks())
+    self.assertEqual(1, self.execute_tasks(status=500))
 
     # Bot retries bot_update, now PubSub works and notification is sent.
     self.publish_successful = True
     self.assertEqual(
         task_result.State.COMPLETED,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='Foo1',
             output_chunk_start=0,
             exit_code=0,
             duration=0.1,
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
     self.assertEqual(0, self.execute_tasks())
-    self.assertEqual(2, len(pub_sub_calls)) # notification is sent
+    self.assertEqual(1, len(pub_sub_calls)) # notification is sent
 
   def _bot_update_timeouts(self, hard, io):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(
         task_result.State.TIMED_OUT,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
             bot_id='localhost',
             cipd_pins=None,
             output='hi',
             output_chunk_start=0,
@@ skipping 66 matching lines @@
 
   def test_bot_update_hard_timeout(self):
     self._bot_update_timeouts(True, False)
 
   def test_bot_update_io_timeout(self):
     self._bot_update_timeouts(False, True)
 
   def test_bot_kill_task(self):
     pub_sub_calls = self.mock_pub_sub()
     dimensions = {u'os': u'Windows-3.1.1', u'pool': u'default'}
-    request = _gen_request(
+    request = self._gen_request(
         properties={'dimensions': dimensions},
         pubsub_topic='projects/abc/topics/def')
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # PENDING -> RUNNING
 
     self.assertEqual(
         None, task_scheduler.bot_kill_task(run_result.key, 'localhost'))
     expected = {
       'abandoned_ts': self.now,
       'bot_dimensions': bot_dimensions,
@@ skipping 49 matching lines @@
       'server_versions': [u'v1a'],
       'started_ts': self.now,
       'state': State.BOT_DIED,
       'try_number': 1,
     }
     self.assertEqual(expected, run_result.key.get().to_dict())
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(2, len(pub_sub_calls)) # RUNNING -> BOT_DIED
 
   def test_bot_kill_task_wrong_bot(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     expected = (
       'Bot bot1 sent task kill for task 1d69b9f088008911 owned by bot '
       'localhost')
     self.assertEqual(
         expected, task_scheduler.bot_kill_task(run_result.key, 'bot1'))
 
   def test_cancel_task(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         },
         pubsub_topic='projects/abc/topics/def')
     pub_sub_calls = self.mock_pub_sub()
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     ok, was_running = task_scheduler.cancel_task(request, result_summary.key)
     self.assertEqual(True, ok)
     self.assertEqual(False, was_running)
     result_summary = result_summary.key.get()
     self.assertEqual(task_result.State.CANCELED, result_summary.state)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # sent completion notification
 
   def test_cancel_task_running(self):
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         },
         pubsub_topic='projects/abc/topics/def')
     pub_sub_calls = self.mock_pub_sub()
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     ok, was_running = task_scheduler.cancel_task(request, result_summary.key)
     self.assertEqual(False, ok)
     self.assertEqual(True, was_running)
     result_summary = result_summary.key.get()
     self.assertEqual(task_result.State.RUNNING, result_summary.state)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # PENDING -> RUNNING
 
   def test_cron_abort_expired_task_to_run(self):
-    request = _gen_request(
+    request = self._gen_request(
             properties={
               'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
             },
             pubsub_topic='projects/abc/topics/def')
     task_request.init_new_request(request, True, None)
     pub_sub_calls = self.mock_pub_sub()
     result_summary = task_scheduler.schedule_request(request, None)
     abandoned_ts = self.mock_now(self.now, request.expiration_secs+1)
     self.assertEqual(
         ['1d69b9f088008910'],
@@ skipping 34 matching lines @@
       'try_number': None,
       'user': u'Jesus',
     }
     self.assertEqual(expected, result_summary.key.get().to_dict())
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # pubsub completion notification
 
   def test_cron_abort_expired_task_to_run_retry(self):
     pub_sub_calls = self.mock_pub_sub()
     now = utils.utcnow()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         },
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=600),
         pubsub_topic='projects/abc/topics/def')
     task_request.init_new_request(request, True, None)
     result_summary = task_scheduler.schedule_request(request, None)
 
     # Fake first try bot died.
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # PENDING -> RUNNING
     now_1 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 1)
     self.assertEqual(([], 1, 0), task_scheduler.cron_handle_bot_died('f.local'))
     self.assertEqual(task_result.State.BOT_DIED, run_result.key.get().state)
     self.assertEqual(
         task_result.State.PENDING, run_result.result_summary_key.get().state)
     self.assertEqual(1, self.execute_tasks())
@@ skipping 43 matching lines @@
     self.assertEqual(expected, result_summary.key.get().to_dict())
 
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(3, len(pub_sub_calls)) # PENDING -> BOT_DIED
 
   def test_cron_handle_bot_died(self):
     pub_sub_calls = self.mock_pub_sub()
 
     # Test first retry, then success.
     now = utils.utcnow()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         },
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=600),
         pubsub_topic='projects/abc/topics/def')
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     self.assertEqual(0, len(pub_sub_calls))
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(
         task_result.State.RUNNING, run_result.result_summary_key.get().state)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # PENDING -> RUNNING
     self.assertEqual(1, run_result.try_number)
     self.assertEqual(task_result.State.RUNNING, run_result.state)
     now_1 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 1)
     self.assertEqual(([], 1, 0), task_scheduler.cron_handle_bot_died('f.local'))
@@ skipping 57 matching lines @@
       ],
       'try_number': 1,
       'user': u'Jesus',
     }
     self.assertEqual(expected, run_result.result_summary_key.get().to_dict())
 
     # Task was retried.
     now_2 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 2)
     bot_dimensions_second = bot_dimensions.copy()
     bot_dimensions_second[u'id'] = [u'localhost-second']
-    _register_bot(bot_dimensions_second)
+    self._register_bot(bot_dimensions_second)
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions_second, 'abc', None)
     self.assertEqual(1, self.execute_tasks())
     self.assertEqual(3, len(pub_sub_calls)) # PENDING -> RUNNING
     logging.info('%s', [t.to_dict() for t in task_to_run.TaskToRun.query()])
     self.assertEqual(2, run_result.try_number)
     self.assertEqual(
         task_result.State.COMPLETED,
         task_scheduler.bot_update_task(
             run_result_key=run_result.key,
@@ skipping 45 matching lines @@
     }
     self.assertEqual(expected, run_result.result_summary_key.get().to_dict())
     self.assertEqual(0.1, run_result.key.get().cost_usd)
 
     self.assertEqual(0, self.execute_tasks())
     self.assertEqual(4, len(pub_sub_calls)) # RUNNING -> COMPLETED
 
   def test_cron_handle_bot_died_same_bot_denied(self):
     # Test first retry, then success.
     now = utils.utcnow()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         },
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=600))
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(1, run_result.try_number)
     self.assertEqual(task_result.State.RUNNING, run_result.state)
     now_1 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 1)
     self.assertEqual(([], 1, 0), task_scheduler.cron_handle_bot_died('f.local'))
 
     # Refresh and compare:
     expected = {
       'abandoned_ts': now_1,
@@ skipping 58 matching lines @@
     now_2 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 2)
     request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(None, request)
     self.assertEqual(None, run_result)
     logging.info('%s', [t.to_dict() for t in task_to_run.TaskToRun.query()])
 
   def test_cron_handle_bot_died_second(self):
     # Test two tries internal_failure's leading to a BOT_DIED status.
     now = utils.utcnow()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
           'idempotent': True,
         },
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=600))
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(1, run_result.try_number)
     self.assertEqual(task_result.State.RUNNING, run_result.state)
     self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 1)
     self.assertEqual(([], 1, 0), task_scheduler.cron_handle_bot_died('f.local'))
     now_1 = self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 2)
     # It must be a different bot.
     bot_dimensions_second = bot_dimensions.copy()
     bot_dimensions_second[u'id'] = [u'localhost-second']
-    _register_bot(bot_dimensions_second)
+    self._register_bot(bot_dimensions_second)
+    # No task to run because the task dimensions were already seen.
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions_second, 'abc', None)
     now_2 = self.mock_now(self.now + 2 * task_result.BOT_PING_TOLERANCE, 3)
     self.assertEqual(
         (['1d69b9f088008912'], 0, 0),
         task_scheduler.cron_handle_bot_died('f.local'))
     self.assertEqual(([], 0, 0), task_scheduler.cron_handle_bot_died('f.local'))
     expected = {
       'abandoned_ts': now_2,
       'bot_dimensions': bot_dimensions_second,
@@ skipping 26 matching lines @@
         u'tag:1',
         u'user:Jesus',
       ],
       'try_number': 2,
       'user': u'Jesus',
     }
     self.assertEqual(expected, run_result.result_summary_key.get().to_dict())
 
   def test_cron_handle_bot_died_ignored_expired(self):
     now = utils.utcnow()
-    request = _gen_request(
+    request = self._gen_request(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         },
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=600))
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
-    _register_bot(bot_dimensions)
+    self._register_bot(bot_dimensions)
     _request, _, run_result = task_scheduler.bot_reap_task(
         bot_dimensions, 'abc', None)
     self.assertEqual(1, run_result.try_number)
     self.assertEqual(task_result.State.RUNNING, run_result.state)
     self.mock_now(self.now + task_result.BOT_PING_TOLERANCE, 601)
     self.assertEqual(
         (['1d69b9f088008911'], 0, 0),
         task_scheduler.cron_handle_bot_died('f.local'))
 
 
 if __name__ == '__main__':
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.CRITICAL)
   unittest.main()