task_queues: Add more scaffolding

appengine/swarming/server/task_to_run_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 hashlib
 import logging
 import os
 import random
@@ skipping 37 matching lines @@
     'priority': 50,
     'properties': task_request.TaskProperties(**props),
     'expiration_ts': now + datetime.timedelta(seconds=60),
     'tags': [u'tag:1'],
     'user': 'Jesus',
   }
   args.update(kwargs)
   return task_request.TaskRequest(**args)
 
 
-def mkreq(req):
-  # This function fits the old style where TaskRequest was stored first, before
-  # TaskToRun and TaskResultSummary.
-  task_request.init_new_request(req, True, None)
-  req.key = task_request.new_request_key()
-  req.put()
-  return req
-
-
 def _task_to_run_to_dict(i):
   """Converts the queue_number to hex for easier testing."""
   out = i.to_dict()
   # Consistent formatting makes it easier to reason about.
   out['queue_number'] = '0x%016x' % out['queue_number']
   return out
 
 
 def _yield_next_available_task_to_dispatch(bot_dimensions, deadline):
   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)
   return [
     _task_to_run_to_dict(to_run)
     for _request, to_run in
         task_to_run.yield_next_available_task_to_dispatch(
             bot_dimensions, deadline)
   ]
 
 
-def _gen_new_task_to_run(**kwargs):
-  """Returns a TaskToRun saved in the DB."""
-  request = mkreq(_gen_request(**kwargs))
-  to_run = task_to_run.new_task_to_run(request)
-  to_run.put()
-  task_queues.assert_task(request)
-  return to_run
-
-
 def _hash_dimensions(dimensions):
   return task_to_run._hash_dimensions(utils.encode_to_json(dimensions))
 
 
 class TestCase(test_case.TestCase):
   def setUp(self):
     super(TestCase, self).setUp()
     auth_testing.mock_get_current_identity(self)
 
 
@@ skipping 160 matching lines @@
 
 
 class TaskToRunApiTest(TestCase):
   def setUp(self):
     super(TaskToRunApiTest, self).setUp()
     self.now = datetime.datetime(2014, 01, 02, 03, 04, 05, 06)
     self.mock_now(self.now)
     # The default expiration_secs for _gen_request().
     self.expiration_ts = self.now + datetime.timedelta(seconds=60)
 
+  def mkreq(self, req, nb_tasks=0):
+    """Stores a new initialized TaskRequest.
+
+    nb_task is 1 or 0. It is 1 when the request.properties.dimensions was new
+    (unseen before) and 0 otherwise.
+    """
+    task_request.init_new_request(req, True, None)
+    task_queues.assert_task(req)
+    self.assertEqual(nb_tasks, self.execute_tasks())
+    req.key = task_request.new_request_key()
+    req.put()
+    return req
+
+  def _gen_new_task_to_run(self, **kwargs):
+    """Returns a TaskToRun saved in the DB."""
+    request = self.mkreq(_gen_request(**kwargs))
+    to_run = task_to_run.new_task_to_run(request)
+    to_run.put()
+    return to_run
+
   def test_all_apis_are_tested(self):
     actual = frozenset(i[5:] for i in dir(self) if i.startswith('test_'))
     # Contains the list of all public APIs.
     expected = frozenset(
         i for i in dir(task_to_run)
         if i[0] != '_' and hasattr(getattr(task_to_run, i), 'func_name'))
     missing = expected - actual
     self.assertFalse(missing)
 
   def test_task_to_run_key_to_request_key(self):
-    request = mkreq(_gen_request())
+    request = self.mkreq(_gen_request())
     task_key = task_to_run.request_to_task_to_run_key(request)
     actual = task_to_run.task_to_run_key_to_request_key(task_key)
     self.assertEqual(request.key, actual)
 
   def test_request_to_task_to_run_key(self):
     self.mock(random, 'getrandbits', lambda _: 0x88)
-    request = mkreq(_gen_request())
+    request = self.mkreq(_gen_request())
     # Ensures that the hash value is constant for the same input.
     self.assertEqual(
         ndb.Key('TaskRequest', 0x7e296460f77ff77e, 'TaskToRun', 3420117132),
         task_to_run.request_to_task_to_run_key(request))
 
   def test_validate_to_run_key(self):
-    request = mkreq(_gen_request())
+    request = self.mkreq(_gen_request())
     task_key = task_to_run.request_to_task_to_run_key(request)
     task_to_run.validate_to_run_key(task_key)
     with self.assertRaises(ValueError):
       task_to_run.validate_to_run_key(ndb.Key('TaskRequest', 1, 'TaskToRun', 1))
 
   def test_gen_queue_number(self):
     # tuples of (input, expected).
     # 2**47 / 365 / 24 / 60 / 60 / 1000. = 4462.756
     data = [
       (('1970-01-01 00:00:00.000',   0), '0x0000000000000000'),
@@ skipping 24 matching lines @@
     data = _gen_request(
         properties={
           'command': [u'command1', u'arg1'],
           'dimensions': request_dimensions,
           'env': {u'foo': u'bar'},
           'execution_timeout_secs': 30,
         },
         priority=20,
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=31))
-    task_to_run.new_task_to_run(mkreq(data)).put()
+    task_to_run.new_task_to_run(self.mkreq(data)).put()
 
     # Create a second with higher priority.
     self.mock(random, 'getrandbits', lambda _: 0x23)
     data = _gen_request(
         properties={
           'command': [u'command1', u'arg1'],
           'dimensions': request_dimensions,
           'env': {u'foo': u'bar'},
           'execution_timeout_secs': 30,
         },
         priority=10,
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=31))
-    task_to_run.new_task_to_run(mkreq(data)).put()
+    task_to_run.new_task_to_run(self.mkreq(data, nb_tasks=0)).put()
 
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.now + datetime.timedelta(seconds=31),
         'request_key': '0x7e296460f77ffdce',
         # Lower priority value means higher priority.
         'queue_number': '0x00060dc5849f1346',
       },
       {
@@ skipping 42 matching lines @@
 
     data_false = (
       ({'os': 'amiga'}, {'os': ['Win', 'Win-3.1']}),
     )
     for request_dimensions, bot_dimensions in data_false:
       self.assertEqual(
           False,
           task_to_run.match_dimensions(request_dimensions, bot_dimensions))
 
   def test_yield_next_available_task_to_dispatch_none(self):
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     # Bot declares no dimensions, so it will fail to match.
     bot_dimensions = {u'id': [u'bot1'], u'pool': [u'default']}
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     self.assertEqual([], actual)
 
   def test_yield_next_available_task_to_dispatch_none_mismatch(self):
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     # Bot declares other dimensions, so it will fail to match.
     bot_dimensions = {
       u'id': [u'bot1'],
       u'os': [u'Windows-3.0'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     self.assertEqual([], actual)
 
   def test_yield_next_available_task_to_dispatch(self):
     request_dimensions = {
       u'foo': u'bar',
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties=dict(dimensions=request_dimensions))
     # Bot declares exactly same dimensions so it matches.
     bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
     bot_dimensions[u'id'] = [u'bot1']
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_dispatch_subset(self):
     request_dimensions = {
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties=dict(dimensions=request_dimensions))
     # Bot declares more dimensions than needed, this is fine and it matches.
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_shard(self):
     request_dimensions = {
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions))
     bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
     bot_dimensions[u'id'] = [u'bot1']
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_dispatch_subset_multivalue(self):
     request_dimensions = {
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties=dict(dimensions=request_dimensions))
     # Bot declares more dimensions than needed.
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows', u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_dispatch_multi_normal(self):
     # Task added one after the other, normal case.
     request_dimensions_1 = {
       u'foo': u'bar',
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
 
     # It's normally time ordered.
     self.mock_now(self.now, 1)
     request_dimensions_2 = {u'id': u'localhost', u'pool': u'default'}
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions_2))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_2))
 
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
@@ skipping 12 matching lines @@
   def test_yield_next_available_task_to_dispatch_clock_skew(self):
     # Asserts that a TaskToRun added later in the DB (with a Key with an higher
     # value) but with a timestamp sooner (for example, time desynchronization
     # between machines) is still returned in the timestamp order, e.g. priority
     # is done based on timestamps and priority only.
     request_dimensions_1 = {
       u'foo': u'bar',
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
 
     # The second shard is added before the first, potentially because of a
     # desynchronized clock. It'll have higher priority.
     self.mock_now(self.now, -1)
     request_dimensions_2 = {u'id': u'localhost', u'pool': u'default'}
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions_2))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_2))
 
     bot_dimensions = {
       u'foo': [u'bar'],
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_2),
         # Due to time being late on the second requester frontend.
         'expiration_ts': self.expiration_ts - datetime.timedelta(seconds=1),
         'queue_number': '0x000a890b67aad106',
       },
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_1),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_dispatch_priority(self):
     # Task added later but with higher priority are returned first.
     request_dimensions_1 = {u'os': u'Windows-3.1.1', u'pool': u'default'}
-    _gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
 
     # This one is later but has higher priority.
     self.mock_now(self.now, 60)
     request_dimensions_2 = {u'os': u'Windows-3.1.1', u'pool': u'default'}
-    _gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions_2), priority=10)
+    request = self.mkreq(
+        _gen_request(
+            properties=dict(dimensions=request_dimensions_2), priority=10),
+        nb_tasks=0)
+    task_to_run.new_task_to_run(request).put()
 
     # It should return them all, in the expected order.
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_1),
         'expiration_ts': datetime.datetime(2014, 1, 2, 3, 6, 5, 6),
         'queue_number': '0x00060dc588329a46',
       },
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_2),
         'expiration_ts': datetime.datetime(2014, 1, 2, 3, 5, 5, 6),
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, None)
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_run_task_exceeds_deadline(self):
     request_dimensions = {
       u'foo': u'bar',
       u'id': u'localhost',
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties=dict(dimensions=request_dimensions))
     # Bot declares exactly same dimensions so it matches.
     bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
     actual = _yield_next_available_task_to_dispatch(
         bot_dimensions, datetime.datetime(1969, 1, 1))
     self.failIf(actual)
 
   def test_yield_next_available_task_to_run_task_meets_deadline(self):
     request_dimensions = {
       u'foo': u'bar',
       u'id': u'localhost',
       u'os': u'Windows-3.1.1',
       u'pool': u'default',
     }
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         properties=dict(dimensions=request_dimensions))
     # Bot declares exactly same dimensions so it matches.
     bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
     actual = _yield_next_available_task_to_dispatch(
         bot_dimensions, datetime.datetime(3000, 1, 1))
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x000a890b67ba1346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_run_task_terminate(self):
     request_dimensions = {
       u'id': u'fake-id',
     }
-    task = _gen_new_task_to_run(
+    task = self._gen_new_task_to_run(
         priority=0,
         properties=dict(
             command=[], dimensions=request_dimensions, execution_timeout_secs=0,
             grace_period_secs=0))
     self.assertTrue(task.key.parent().get().properties.is_terminate)
     # Bot declares exactly same dimensions so it matches.
     bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
     bot_dimensions[u'pool'] = [u'default']
     actual = _yield_next_available_task_to_dispatch(bot_dimensions, 0)
     expected = [
       {
         'dimensions_hash': _hash_dimensions(request_dimensions),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x0004eef40bd85346',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_expired_task_to_run(self):
     now = utils.utcnow()
-    _gen_new_task_to_run(
+    self._gen_new_task_to_run(
         created_ts=now,
         expiration_ts=now+datetime.timedelta(seconds=60))
     bot_dimensions = {u'id': [u'bot1'], u'pool': [u'default']}
     self.assertEqual(
         1,
         len(_yield_next_available_task_to_dispatch(bot_dimensions, None)))
     self.assertEqual(
         0, len(list(task_to_run.yield_expired_task_to_run())))
 
     # All tasks are now expired. Note that even if they still have .queue_number
     # set because the cron job wasn't run, they are still not yielded by
     # yield_next_available_task_to_dispatch()
     self.mock_now(self.now, 61)
     self.assertEqual(
         0, len(_yield_next_available_task_to_dispatch(bot_dimensions, None)))
     self.assertEqual(
         1, len(list(task_to_run.yield_expired_task_to_run())))
 
   def test_is_reapable(self):
     req_dimensions = {u'os': u'Windows-3.1.1', u'pool': u'default'}
-    to_run = _gen_new_task_to_run(properties=dict(dimensions=req_dimensions))
+    to_run = self._gen_new_task_to_run(
+        properties=dict(dimensions=req_dimensions))
     bot_dimensions = {
       u'id': [u'localhost'],
       u'os': [u'Windows-3.1.1'],
       u'pool': [u'default'],
     }
     self.assertEqual(
         1, len(_yield_next_available_task_to_dispatch(bot_dimensions, None)))
 
     self.assertEqual(True, to_run.is_reapable)
     to_run.queue_number = None
     to_run.put()
     self.assertEqual(False, to_run.is_reapable)
 
   def test_set_lookup_cache(self):
-    to_run = _gen_new_task_to_run(
+    to_run = self._gen_new_task_to_run(
         properties={
           'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
         })
     self.assertEqual(False, task_to_run._lookup_cache_is_taken(to_run.key))
     task_to_run.set_lookup_cache(to_run.key, True)
     self.assertEqual(False, task_to_run._lookup_cache_is_taken(to_run.key))
     task_to_run.set_lookup_cache(to_run.key, False)
     self.assertEqual(True, task_to_run._lookup_cache_is_taken(to_run.key))
     task_to_run.set_lookup_cache(to_run.key, True)
     self.assertEqual(False, task_to_run._lookup_cache_is_taken(to_run.key))
 
 
 if __name__ == '__main__':
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   logging.basicConfig(
       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
   unittest.main()