Add support for repeated keys in TaskRequest.

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 88 matching lines @@
       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'},
+      'dimensions_flat': [u'pool:default'],
       'env': {},
       'execution_timeout_secs': 24*60*60,
       'io_timeout_secs': None,
     }
     props.update(properties or {})
-    props['dimensions_dict'] = props.pop('dimensions')
     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=1):
     """Schedules a task.
 
     nb_task is 1 if a GAE task queue rebuild-task-cache was enqueued.
     """
-    request = self._gen_request(properties={'dimensions': dims})
+    self.assertIsInstance(dims, list)
+    request = self._gen_request(properties={'dimensions_flat': 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=1):
     """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=1):
     """Reaps a task."""
-    self._quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'})
+    self._quick_schedule([u'os:Windows-3.1.1', u'pool: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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     self._register_bot(bot_dimensions)
@@ skipping 43 matching lines @@
       'auth_token': 'token',
       'userdata': 'userdata',
     })
     self.assertEqual([
       {
         'attributes': {'auth_token': 'token'},
         'message': '{"task_id":"abcdef123","userdata":"userdata"}',
         'topic': 'projects/abc/topics/def',
     }], calls)
 
-  def _task_ran_successfully(self, nb_task=1):
+  def _task_ran_successfully(self, dimensions_flat=None, nb_task=1):
     """Runs a task successfully and returns the task_id."""
+    dimensions_flat = dimensions_flat or [u'os:Windows-3.1.1', u'pool:default']
     request = self._gen_request(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-          'idempotent': True,
-        })
+        properties={'dimensions_flat': dimensions_flat, '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'],
+      u'pool': [u'default', u'testing'],
     }
     self._register_bot(bot_dimensions, nb_task=nb_task)
     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,
@@ skipping 10 matching lines @@
             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_task=1, now=None):
     request = self._gen_request(
         name='yay',
         user='Raoul',
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
+      u'pool': [u'default', u'testing'],
     }
     self.assertEqual(None, task_to_run.TaskToRun.query().get().queue_number)
     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)
     result_summary_duped, run_results_duped = get_results(request.key)
     expected = {
       'abandoned_ts': None,
       'bot_dimensions': bot_dimensions,
@@ skipping 28 matching lines @@
         u'service_account:none',
         u'tag:1',
         u'user:Raoul',
       ],
       'try_number': 0,
       'user': u'Raoul',
     }
     self.assertEqual(expected, result_summary_duped.to_dict())
     self.assertEqual([], run_results_duped)
 
+  def test_task_two_pools(self):
+    self._task_ran_successfully(
+        dimensions_flat=[u'pool:default', u'pool:testing'])
+
   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')
 
   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 = self._gen_request(
         name='yay',
         user='Raoul',
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool:default'],
           'idempotent': True,
         })
     task_request.init_new_request(request, True, None)
     _result_summary = task_scheduler.schedule_request(request, None)
     self.assertEqual(1, self.execute_tasks())
     # 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')
 
     # 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 = self._gen_request(
         name='yay',
         user='Jesus',
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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_variable(self):
     # Test the edge case where config.settings().reusable_task_age_secs is being
     # modified. This ensure TaskResultSummary.order(TRS.key) works.
@@ skipping 17 matching lines @@
     third_ts = self.mock_now(self.now, 20)
     self._task_deduped(
         third_ts, task_id, '1d69ba3ea8008b10', nb_task=0, now=second_ts)
 
   def test_task_parent_children(self):
     # Parent task creates a child task.
     parent_id = self._task_ran_successfully()
     request = self._gen_request(
         parent_task_id=parent_id,
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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 = self._gen_request(
         properties={
           'command': [],
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
@@ skipping 21 matching lines @@
             hard_timeout=False,
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
 
     parent_id = run_result.task_id
     request = self._gen_request(
         parent_task_id=parent_id,
         properties={
-          'dimensions':{u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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 = {
       'abandoned_ts': None,
       'bot_dimensions': None,
@@ skipping 222 matching lines @@
         '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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
@@ skipping 74 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(
-        self._quick_schedule({u'os': u'Windows-3.1.1', u'pool': u'default'}))
+        self._quick_schedule([u'os:Windows-3.1.1', u'pool: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'})
-    self._quick_schedule({u'pool': u'good'})
+    self._quick_schedule([u'pool:good'])
     with self.assertRaises(auth.AuthorizationError):
-      self._quick_schedule({u'pool': u'bad'})
+      self._quick_schedule([u'pool:bad'])
 
   def test_schedule_request_forbidden_dim_via_star(self):
     self.mock_dim_acls({u'abc:*': u'noone'})
-    self._quick_schedule({u'pool': u'default'})
+    self._quick_schedule([u'pool:default'])
     with self.assertRaises(auth.AuthorizationError):
-      self._quick_schedule({u'pool': u'default', u'abc': u'blah'})
+      self._quick_schedule([u'abc:blah', u'pool:default'])
 
   def test_schedule_request_id_without_pool(self):
     self.mock_dim_acls({u'pool:good': u'mocked'})
     with self.assertRaises(auth.AuthorizationError):
-      self._quick_schedule({u'id': u'abc'})
+      self._quick_schedule([u'id:abc'])
     auth_testing.mock_is_admin(self)
-    self._quick_schedule({u'id': u'abc'}, nb_task=0)
+    self._quick_schedule([u'id:abc'], nb_task=0)
 
   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)
 
-    self._quick_schedule({u'id': u'abc', u'pool': u'unknown'}, nb_task=0)
-    self._quick_schedule({u'id': u'abc', u'pool': u'good'}, nb_task=0)
+    self._quick_schedule([u'id:abc', u'pool:unknown'], nb_task=0)
+    self._quick_schedule([u'id:abc', u'pool:good'], nb_task=0)
     with self.assertRaises(auth.AuthorizationError):
-      self._quick_schedule({u'id': u'abc', u'pool': u'bad'})
+      self._quick_schedule([u'id:abc', u'pool:bad'])
 
   def test_bot_update_task(self):
     run_result = self._quick_reap(nb_task=0)
     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',
@@ skipping 75 matching lines @@
             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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
@@ skipping 37 matching lines @@
             io_timeout=False,
             cost_usd=0.1,
             outputs_ref=None,
             performance_stats=None))
     self.assertEqual(0, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # notification is sent
 
   def _bot_update_timeouts(self, hard, io):
     request = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
@@ skipping 74 matching lines @@
     self.assertEqual(expected, run_result.key.get().to_dict())
 
   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 = self._gen_request(
-        properties={'dimensions': dimensions},
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool:default']},
         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'],
     }
     self._register_bot(bot_dimensions)
     reaped_request, _, run_result = task_scheduler.bot_reap_task(
@@ skipping 59 matching lines @@
       '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 = self._gen_request(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        })
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     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 = self._gen_request(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        },
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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(2, self.execute_tasks())
     self.assertEqual(1, len(pub_sub_calls)) # sent completion notification
 
   def test_cancel_task_running(self):
     request = self._gen_request(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        },
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     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 = self._gen_request(
-            properties={
-              'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-            },
-            pubsub_topic='projects/abc/topics/def')
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
         task_scheduler.cron_abort_expired_task_to_run('f.local'))
     self.assertEqual([], task_result.TaskRunResult.query().fetch())
     expected = {
       'abandoned_ts': abandoned_ts,
@@ skipping 32 matching lines @@
     }
     self.assertEqual(expected, result_summary.key.get().to_dict())
     self.assertEqual(2, 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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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 = {
@@ skipping 61 matching lines @@
     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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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(1, self.execute_tasks())
     self.assertEqual(0, len(pub_sub_calls))
     bot_dimensions = {
@@ skipping 141 matching lines @@
     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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
@@ skipping 73 matching lines @@
         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 = self._gen_request(
         properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
+          'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
@@ skipping 52 matching lines @@
         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 = self._gen_request(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        },
+        properties={'dimensions_flat': [u'os:Windows-3.1.1', u'pool: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'],
     }
     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()