swarming: mechanical changes to unit tests

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
 import sys
 import timeit
 import unittest
 
-import test_env
-test_env.setup_test_env()
+# Setups environment.
+APP_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, APP_DIR)
+import test_env_handlers
+
+import webtest
 
 from google.appengine.ext import ndb
 
+import handlers_backend
+
 from components import auth_testing
 from components import utils
 from test_support import test_case
 
 from server import bot_management
 from server import task_queues
 from server import task_request
 from server import task_to_run
 
 
@@ skipping 43 matching lines @@
     for _request, to_run in
         task_to_run.yield_next_available_task_to_dispatch(
             bot_dimensions, deadline)
   ]
 
 
 def _hash_dimensions(dimensions):
   return task_to_run._hash_dimensions(utils.encode_to_json(dimensions))
 
 
-class TestCase(test_case.TestCase):
+class TaskToRunPrivateTest(test_case.TestCase):
   def setUp(self):
-    super(TestCase, self).setUp()
+    super(TaskToRunPrivateTest, self).setUp()
     auth_testing.mock_get_current_identity(self)
 
-
-class TaskToRunPrivateTest(TestCase):
   def test_powerset(self):
     # tuples of (input, expected).
     # TODO(maruel): We'd want the code to deterministically try 'Windows-6.1'
     # before 'Windows'. Probably do a reverse() on the values?
     data = [
       ({'OS': 'Windows'}, [{'OS': 'Windows'}, {}]),
       (
         {'OS': ['Windows', 'Windows-6.1']},
         [{'OS': 'Windows'}, {'OS': 'Windows-6.1'}, {}],
       ),
@@ skipping 141 matching lines @@
     self.assertEqual(16384, len(items))
 
   def test_dimensions_search_sizing_14_1(self):
     dimensions = {str(k): '01234567890123456789' for k in xrange(14)}
     items = tuple(sorted(
         task_to_run._hash_dimensions(utils.encode_to_json(i))
         for i in task_to_run._powerset(dimensions)))
     self.assertEqual(16384, len(items))
 
 
-class TaskToRunApiTest(TestCase):
+class TaskToRunApiTest(test_env_handlers.AppTestBase):
   def setUp(self):
     super(TaskToRunApiTest, self).setUp()
     self.now = datetime.datetime(2014, 01, 02, 03, 04, 05, 06)
     self.mock_now(self.now)
+    auth_testing.mock_get_current_identity(self)
     # The default expiration_secs for _gen_request().
     self.expiration_ts = self.now + datetime.timedelta(seconds=60)
+    # Setup the backend to handle task queues for 'task-dimensions'.
+    self.app = webtest.TestApp(
+        handlers_backend.create_application(True),
+        extra_environ={
+          'REMOTE_ADDR': self.source_ip,
+          'SERVER_SOFTWARE': os.environ['SERVER_SOFTWARE'],
+        })
+    self._enqueue_orig = self.mock(utils, 'enqueue_task', self._enqueue)
 
-  def mkreq(self, req, nb_tasks=0):
+  def _enqueue(self, *args, **kwargs):
+    return self._enqueue_orig(*args, use_dedicated_module=False, **kwargs)
+
+  def mkreq(self, req, nb_task=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())
+    self.assertEqual(nb_task, 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
@@ skipping 74 matching lines @@
     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(self.mkreq(data, nb_tasks=0)).put()
+    task_to_run.new_task_to_run(self.mkreq(data, nb_task=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 236 matching lines @@
     # Task added later but with higher priority are returned first.
     request_dimensions_1 = {u'os': u'Windows-3.1.1', u'pool': u'default'}
     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'}
     request = self.mkreq(
         _gen_request(
             properties=dict(dimensions=request_dimensions_2), priority=10),
-        nb_tasks=0)
+        nb_task=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',
       },
       {
@@ skipping 121 matching lines @@
     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()