Add support for repeated keys in TaskRequest.

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 logging
 import os
 import random
 import sys
@@ skipping 17 matching lines @@
 from server import bot_management
 from server import task_queues
 from server import task_request
 from server import task_to_run
 
 
 # pylint: disable=W0212
 # Method could be a function - pylint: disable=R0201
 
 
+def _flatten(request_dimensions):
+  return [u'%s:%s' % (k, v) for k, v in request_dimensions]
+
+
 def _gen_request(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',
   }
@@ skipping 123 matching lines @@
     ]
     for i, ((dimensions_hash, timestamp, priority), expected) in enumerate(
         data):
       d = datetime.datetime.strptime(timestamp, '%Y-%m-%d %H:%M:%S.%f')
       actual = '0x%016x' % task_to_run._gen_queue_number(
           dimensions_hash, d, priority)
       self.assertEquals((i, expected), (i, actual))
 
   def test_new_task_to_run(self):
     self.mock(random, 'getrandbits', lambda _: 0x12)
-    request_dimensions = {u'os': u'Windows-3.1.1', u'pool': u'default'}
+    request_dimensions = [(u'os', u'Windows-3.1.1'), (u'pool', u'default')]
     now = utils.utcnow()
     data = _gen_request(
         properties={
           'command': [u'command1', u'arg1'],
-          'dimensions': request_dimensions,
+          'dimensions_flat': _flatten(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(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,
+          'dimensions_flat': _flatten(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_task=0)).put()
 
     expected = [
       {
@@ skipping 17 matching lines @@
       return out
 
     # Warning: Ordering by key doesn't work because of TaskToRunShard; e.g.
     # the entity key ordering DOES NOT correlate with .queue_number
     # Ensure they come out in expected order.
     q = task_to_run.TaskToRun.query().order(task_to_run.TaskToRun.queue_number)
     self.assertEqual(expected, map(flatten, q.fetch()))
 
   def test_match_dimensions(self):
     data_true = (
-      ({}, {}),
-      ({}, {'a': 'b'}),
-      ({'a': 'b'}, {'a': 'b'}),
-      ({'os': 'amiga'}, {'os': ['amiga', 'amiga-3.1']}),
-      ( {'os': 'amiga', 'foo': 'bar'},
+      ([], {}),
+      ([], {'a': 'b'}),
+      ([('a', 'b')], {'a': 'b'}),
+      ([('os', 'amiga')], {'os': ['amiga', 'amiga-3.1']}),
+      ( [('os', 'amiga'), ('foo', 'bar')],
         {'os': ['amiga', 'amiga-3.1'], 'a': 'b', 'foo': 'bar'}),
     )
 
     for request_dimensions, bot_dimensions in data_true:
       self.assertEqual(
           True,
           task_to_run.match_dimensions(request_dimensions, bot_dimensions))
 
     data_false = (
-      ({'os': 'amiga'}, {'os': ['Win', 'Win-3.1']}),
+      ([('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):
+    request_dimensions = [
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        })
+        properties=dict(dimensions_flat=_flatten(request_dimensions)))
     # 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):
+    request_dimensions = [
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        })
+        properties=dict(dimensions_flat=_flatten(request_dimensions)))
     # 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',
-    }
+    request_dimensions = [
+      (u'foo', u'bar'),
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions))
+        properties=dict(dimensions_flat=_flatten(request_dimensions)))
     # Bot declares exactly same dimensions so it matches.
-    bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
+    bot_dimensions = {k: [v] for k, v in request_dimensions}
     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': '0x613fbb330c8ede72',
       },
     ]
     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',
-    }
+    request_dimensions = [
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions))
+        properties=dict(dimensions_flat=_flatten(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': '0x1a3aa6630c8ede72',
       },
     ]
     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',
-    }
-    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions))
-    bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
+    request_dimensions = [
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
+    self._gen_new_task_to_run(
+        properties=dict(dimensions_flat=_flatten(request_dimensions)))
+    bot_dimensions = {k: [v] for k, v in request_dimensions}
     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': '0x1a3aa6630c8ede72',
       },
     ]
     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',
-    }
+    request_dimensions = [
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions))
+        properties=dict(dimensions_flat=_flatten(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': '0x1a3aa6630c8ede72',
       },
     ]
     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',
-    }
-    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    request_dimensions_1 = [
+      (u'foo', u'bar'),
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
+    self._gen_new_task_to_run(
+        properties=dict(dimensions_flat=_flatten(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'}
+    request_dimensions_2 = [(u'id', u'localhost'), (u'pool', u'default')]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions_2), nb_task=0)
+        properties=dict(dimensions_flat=_flatten(request_dimensions_2)),
+        nb_task=0)
 
     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_1),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x613fbb330c8ede72',
       },
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_2),
         'expiration_ts': self.expiration_ts + datetime.timedelta(seconds=1),
         'queue_number': '0x5385bf748c8ede7c',
       },
     ]
     # There is a significant risk of non-determinism.
     self.assertEqual(sorted(expected), sorted(actual))
 
   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',
-    }
-    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    request_dimensions_1 = [
+      (u'foo', u'bar'),
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
+    self._gen_new_task_to_run(
+        properties=dict(dimensions_flat=_flatten(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'}
+    request_dimensions_2 = [(u'id', u'localhost'), (u'pool', u'default')]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions_2), nb_task=0)
+        properties=dict(dimensions_flat=_flatten(request_dimensions_2)),
+        nb_task=0)
 
     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_1),
         'expiration_ts': self.expiration_ts,
         'queue_number': '0x613fbb330c8ede72',
       },
       {
         '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': '0x5385bf748c8ede68',
       },
     ]
     # There is a significant risk of non-determinism.
     self.assertEqual(sorted(expected), sorted(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'}
-    self._gen_new_task_to_run(properties=dict(dimensions=request_dimensions_1))
+    request_dimensions_1 = [(u'os', u'Windows-3.1.1'), (u'pool', u'default')]
+    self._gen_new_task_to_run(
+        properties=dict(dimensions_flat=_flatten(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_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),
+            properties=dict(dimensions_flat=_flatten(request_dimensions_2)),
+            priority=10),
         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': '0x1a3aa663028ee0ca',
       },
       {
         'dimensions_hash': _hash_dimensions(request_dimensions_2),
         'expiration_ts': datetime.datetime(2014, 1, 2, 3, 5, 5, 6),
         'queue_number': '0x1a3aa6630c8ede72',
       },
     ]
     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',
-    }
+    request_dimensions = [
+      (u'foo', u'bar'),
+      (u'id', u'localhost'),
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions), nb_task=0)
+        properties={
+          'dimensions_flat': [u'%s:%s' % (k, v) for k, v in request_dimensions],
+        },
+        nb_task=0)
     # Bot declares exactly same dimensions so it matches.
-    bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
+    bot_dimensions = {k: [v] for k, v in request_dimensions}
     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',
-    }
+    request_dimensions = [
+      (u'foo', u'bar'),
+      (u'id', u'localhost'),
+      (u'os', u'Windows-3.1.1'),
+      (u'pool', u'default'),
+    ]
     self._gen_new_task_to_run(
-        properties=dict(dimensions=request_dimensions), nb_task=0)
+        properties=dict(dimensions_flat=_flatten(request_dimensions)),
+        nb_task=0)
     # Bot declares exactly same dimensions so it matches.
-    bot_dimensions = {k: [v] for k, v in request_dimensions.iteritems()}
+    bot_dimensions = {k: [v] for k, v in request_dimensions}
     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': '0x3f6b0f050c8ede72',
       },
     ]
     self.assertEqual(expected, actual)
 
   def test_yield_next_available_task_to_run_task_terminate(self):
-    request_dimensions = {
-      u'id': u'fake-id',
-    }
+    request_dimensions = [
+      (u'id', u'fake-id'),
+    ]
     task = self._gen_new_task_to_run(
         priority=0,
         properties=dict(
-            command=[], dimensions=request_dimensions, execution_timeout_secs=0,
-            grace_period_secs=0),
+            command=[], dimensions_flat=_flatten(request_dimensions),
+            execution_timeout_secs=0, grace_period_secs=0),
         nb_task=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 = {k: [v] for k, v in request_dimensions}
     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': '0x54795e3c800ede72',
       },
     ]
     self.assertEqual(expected, actual)
@@ skipping 13 matching lines @@
     # 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'}
+    req_dimensions = [(u'os', u'Windows-3.1.1'), (u'pool', u'default')]
     to_run = self._gen_new_task_to_run(
-        properties=dict(dimensions=req_dimensions))
+        properties=dict(dimensions_flat=_flatten(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):
+    req_dimensions = [(u'os', u'Windows-3.1.1'), (u'pool', u'default')]
     to_run = self._gen_new_task_to_run(
-        properties={
-          'dimensions': {u'os': u'Windows-3.1.1', u'pool': u'default'},
-        })
+        properties=dict(dimensions_flat=_flatten(req_dimensions)))
     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()