| #!/usr/bin/env vpython |
| # coding: utf-8 |
| # Copyright 2017 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 json |
| import logging |
| import os |
| import random |
| import sys |
| import unittest |
| |
| # Setups environment. |
| APP_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| sys.path.insert(0, APP_DIR) |
| import test_env_handlers |
| |
| from parameterized import parameterized |
| import webtest |
| |
| from google.appengine.api import datastore_errors |
| from google.appengine.api import memcache |
| from google.appengine.ext import ndb |
| |
| import handlers_backend |
| |
| from components import datastore_utils |
| from components import utils |
| from server import bot_management |
| from server import task_queues |
| from server import task_request |
| |
| |
| def _gen_properties(**kwargs): |
| """Creates a TaskProperties.""" |
| args = { |
| "command": ["command1"], |
| "dimensions": { |
| "cpu": ["x86-64"], |
| "os": ["Ubuntu-16.04"], |
| "pool": ["default"], |
| }, |
| "env": {}, |
| "execution_timeout_secs": 24 * 60 * 60, |
| "io_timeout_secs": None, |
| } |
| args.update(kwargs) |
| args["dimensions_data"] = args.pop("dimensions") |
| return task_request.TaskProperties(**args) |
| |
| |
| def _gen_request(properties=None): |
| """Creates a TaskRequest.""" |
| now = utils.utcnow() |
| args = { |
| "created_ts": now, |
| "manual_tags": ["tag:1"], |
| "name": "Request name", |
| "priority": 50, |
| "task_slices": [ |
| task_request.TaskSlice( |
| expiration_secs=60, properties=properties or _gen_properties() |
| ), |
| ], |
| "user": "Jesus", |
| } |
| req = task_request.TaskRequest(**args) |
| task_request.init_new_request(req, True, task_request.TEMPLATE_AUTO) |
| return req |
| |
| |
| def _to_timestamp(dt): |
| return int(utils.datetime_to_timestamp(dt) / 1e6) |
| |
| |
| class TaskQueuesApiTest(test_env_handlers.AppTestBase): |
| def setUp(self): |
| super(TaskQueuesApiTest, self).setUp() |
| # Setup the backend to handle task queues. |
| self.app = webtest.TestApp( |
| handlers_backend.create_application(True), |
| extra_environ={ |
| "REMOTE_ADDR": self.source_ip, |
| "SERVER_SOFTWARE": os.environ["SERVER_SOFTWARE"], |
| }, |
| ) |
| |
| # Flip use_dedicated_module to False, we don't have it mocked in tests. |
| enqueue_task_async = utils.enqueue_task_async |
| |
| def enqueue_task_async_mock(*args, **kwargs): |
| return enqueue_task_async(*args, use_dedicated_module=False, **kwargs) |
| |
| self.mock(utils, "enqueue_task_async", enqueue_task_async_mock) |
| |
| def random_dt(a, b): |
| return datetime.timedelta(minutes=(a + b) / 2.0) |
| |
| self.mock(task_queues, "_random_timedelta_mins", random_dt) |
| |
| def _assert_bot(self, bot_id="bot1", dimensions=None): |
| bot_dimensions = { |
| "cpu": ["x86-64", "x64"], |
| "id": [bot_id], |
| "os": ["Ubuntu-16.04", "Ubuntu"], |
| "pool": ["default"], |
| } |
| bot_dimensions.update(dimensions or {}) |
| bot_management.bot_event( |
| event_type="request_sleep", |
| bot_id=bot_id, |
| external_ip="1.2.3.4", |
| authenticated_as=bot_id, |
| dimensions=bot_dimensions, |
| state={}, |
| version="1234", |
| register_dimensions=True, |
| ) |
| queues = task_queues.assert_bot(bot_dimensions) |
| return self.execute_tasks(), queues |
| |
| def _assert_task(self, dimensions=None): |
| """Creates one pending TaskRequest and asserts it in task_queues.""" |
| request = _gen_request( |
| properties=_gen_properties(dimensions=dimensions) if dimensions else None |
| ) |
| task_queues.assert_task_async(request).get_result() |
| return self.execute_tasks() |
| |
| 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_queues) |
| if i[0] != "_" and hasattr(getattr(task_queues, i), "func_name") |
| ) |
| missing = expected - actual |
| self.assertFalse(missing) |
| |
| def assert_count(self, count, entity): |
| actual = entity.query().count() |
| if actual != count: |
| self.fail([i.to_dict() for i in entity.query()]) |
| self.assertEqual(count, actual) |
| |
| def test_assert_bot(self): |
| self.assert_count(0, task_queues.BotDimensionsMatches) |
| self.assert_count(0, task_queues.TaskDimensionsInfo) |
| self.assert_count(0, task_queues.TaskDimensionsSets) |
| self._assert_bot() |
| self.assert_count(1, task_queues.BotDimensionsMatches) |
| self.assert_count(0, task_queues.TaskDimensionsInfo) |
| self.assert_count(0, task_queues.TaskDimensionsSets) |
| |
| def test_assert_task_async(self): |
| self.assert_count(0, task_queues.BotDimensionsMatches) |
| self.assert_count(0, task_queues.TaskDimensionsInfo) |
| self.assert_count(0, task_queues.TaskDimensionsSets) |
| self._assert_task() |
| self.assert_count(0, task_queues.BotDimensionsMatches) |
| self.assert_count(1, task_queues.TaskDimensionsInfo) |
| self.assert_count(1, task_queues.TaskDimensionsSets) |
| |
| def test_assert_task_async_no_update(self): |
| # Ran TQ tasks to register the new dimension set. |
| tq_tasks = self._assert_task() |
| self.assertEqual(tq_tasks, 1) |
| # Already seen it, no new tasks. |
| tq_tasks = self._assert_task() |
| self.assertEqual(tq_tasks, 0) |
| |
| def test_assert_task_async_or_dims(self): |
| # Ran TQ tasks to register the new dimension set. |
| tq_tasks = self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1|v2"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| self.assertEqual(tq_tasks, 1) |
| # Already seen it, no new tasks. |
| tq_tasks = self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1|v2"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| self.assertEqual(tq_tasks, 0) |
| |
| def test_freshen_up_queues(self): |
| # See more complex test below. |
| pass |
| |
| def test_or_dimensions_new_tasks(self): |
| # There are two bots that can handle two OR alternatives of a task. |
| self._assert_bot( |
| bot_id="bot1", |
| dimensions={ |
| "os": ["v1", "v2"], |
| "gpu": ["nv"], |
| }, |
| ) |
| self._assert_bot( |
| bot_id="bot2", |
| dimensions={ |
| "os": ["v2"], |
| "gpu": ["amd"], |
| }, |
| ) |
| |
| # A task that "|" appears. |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1|v2"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| |
| # Both bots should be able to handle it. |
| self.assertEqual([787294789], task_queues.freshen_up_queues("bot1")) |
| self.assertEqual([787294789], task_queues.freshen_up_queues("bot2")) |
| |
| # Another task comes in that matches only bot1. |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| |
| # Only bot1 can handle the second task. |
| self.assertEqual( |
| [787294789, 1873318153], task_queues.freshen_up_queues("bot1") |
| ) |
| self.assertEqual([787294789], task_queues.freshen_up_queues("bot2")) |
| |
| def test_or_dimensions_new_bots(self): |
| # Register a task that uses "|" dimensions. |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1|v2"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| # Register a task that matches only bot1. |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1"], |
| "gpu": ["nv|amd"], |
| } |
| ) |
| |
| # Two bots show up matching two different OR alternatives of the first task. |
| self._assert_bot( |
| bot_id="bot1", |
| dimensions={ |
| "os": ["v1", "v2"], |
| "gpu": ["nv"], |
| }, |
| ) |
| self._assert_bot( |
| bot_id="bot2", |
| dimensions={ |
| "os": ["v2"], |
| "gpu": ["amd"], |
| }, |
| ) |
| |
| # The bots got correct matches. |
| self.assertEqual( |
| [787294789, 1873318153], task_queues.freshen_up_queues("bot1") |
| ) |
| self.assertEqual([787294789], task_queues.freshen_up_queues("bot2")) |
| |
| def test_or_dimensions_same_hash(self): |
| self._assert_bot(bot_id="bot1", dimensions={"os": ["v1"]}) |
| self._assert_bot(bot_id="bot2", dimensions={"os": ["v2"]}) |
| self._assert_bot(bot_id="bot3", dimensions={"os": ["v3"]}) |
| |
| # Both requests should have the same dimension_hash. |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v1|v2|v3"], |
| } |
| ) |
| self._assert_task( |
| { |
| "pool": ["default"], |
| "os": ["v3|v2|v1"], |
| } |
| ) |
| |
| # All 3 bots are matched to the same queue. |
| self.assertEqual([1060356668], task_queues.freshen_up_queues("bot1")) |
| self.assertEqual([1060356668], task_queues.freshen_up_queues("bot2")) |
| self.assertEqual([1060356668], task_queues.freshen_up_queues("bot3")) |
| |
| def test_bot_dimensions_to_flat(self): |
| actual = task_queues.bot_dimensions_to_flat( |
| {"a": ["c", "bee"], "cee": ["zee"]} |
| ) |
| self.assertEqual(["a:bee", "a:c", "cee:zee"], actual) |
| |
| def test_dimensions_to_flat_long_ascii(self): |
| key = "a" * 64 |
| actual = task_queues.bot_dimensions_to_flat( |
| { |
| key: [ |
| # Too long. |
| "b" * 257, |
| # Ok. |
| "c" * 256, |
| ], |
| } |
| ) |
| expected = [ |
| key + ":" + "b" * 256 + "…", |
| key + ":" + "c" * 256, |
| ] |
| self.assertEqual(expected, actual) |
| |
| def test_dimensions_to_flat_long_unicode(self): |
| key = "a" * 64 |
| actual = task_queues.bot_dimensions_to_flat( |
| { |
| key: [ |
| # Ok. |
| "⌛" * 256, |
| # Too long. |
| "â›”" * 257, |
| ], |
| } |
| ) |
| expected = [ |
| key + ":" + "⌛" * 256, |
| key + ":" + "⛔" * 256 + "…", |
| ] |
| self.assertEqual(expected, actual) |
| |
| def test_dimensions_to_flat_long_unicode_non_BMP(self): |
| # For non-BMP characters, the length is effectively halved for now. |
| key = "a" * 64 |
| # Python considers emoji in the supplemental plane to have length 2 on UCS2 |
| # builds, and length 1 on UCS4 builds. |
| l = 128 if sys.maxunicode < 65536 else 256 |
| actual = task_queues.bot_dimensions_to_flat( |
| { |
| key: [ |
| # Too long. |
| "💥" * (l + 1), |
| # Ok. |
| "😬" * l, |
| ], |
| } |
| ) |
| expected = [ |
| key + ":" + "💥" * l + "…", |
| key + ":" + "😬" * l, |
| ] |
| self.assertEqual(expected, actual) |
| |
| def test_dimensions_to_flat_duplicate_value(self): |
| actual = task_queues.bot_dimensions_to_flat({"a": ["c", "c"]}) |
| self.assertEqual(["a:c"], actual) |
| |
| def test_python_len_non_BMP(self): |
| # Here are emojis in the base plane. They are 1 character. |
| self.assertEqual(1, len("⌛")) |
| self.assertEqual(1, len("â›”")) |
| # Python considers emoji in the supplemental plane to have length 2 on UCS2 |
| # builds, and length 1 on UCS4 builds. |
| l = 2 if sys.maxunicode < 65536 else 1 |
| self.assertEqual(l, len("😬")) |
| self.assertEqual(l, len("💥")) |
| |
| def test_probably_has_capacity_empty(self): |
| # The bot can service this dimensions. |
| d = {"pool": ["default"], "os": ["Ubuntu-16.04"]} |
| # By default, nothing has capacity. |
| self.assertEqual(None, task_queues.probably_has_capacity(d)) |
| |
| def test_probably_has_capacity(self): |
| d = {"pool": ["default"], "os": ["Ubuntu-16.04"]} |
| # A bot coming online doesn't register capacity automatically. |
| self._assert_bot() |
| self.assertEqual(1, bot_management.BotInfo.query().count()) |
| self.assertEqual(None, task_queues.probably_has_capacity(d)) |
| |
| def test_probably_has_capacity_freshen_up_queues(self): |
| d = {"pool": ["default"], "os": ["Ubuntu-16.04"]} |
| # Capacity registers there only once there's a request enqueued and |
| # freshen_up_queues() is called. |
| self._assert_bot() |
| request = _gen_request(properties=_gen_properties(dimensions=d)) |
| task_queues.assert_task_async(request).get_result() |
| self.execute_tasks() |
| self.assertEqual(None, task_queues.probably_has_capacity(d)) |
| |
| # It gets set only once freshen_up_queues() is called. |
| task_queues.freshen_up_queues("bot1") |
| self.assertEqual(True, task_queues.probably_has_capacity(d)) |
| |
| def test_set_has_capacity(self): |
| d = {"pool": ["default"], "os": ["Ubuntu-16.04"]} |
| # By default, nothing has capacity. None means no data. |
| now = utils.utcnow() |
| self.mock_now(now, 0) |
| self.assertEqual(None, task_queues.probably_has_capacity(d)) |
| # Keep the value for 2 seconds, exclusive. |
| task_queues.set_has_capacity(d, 2) |
| self.assertEqual(True, task_queues.probably_has_capacity(d)) |
| self.mock_now(now, 1) |
| self.assertEqual(True, task_queues.probably_has_capacity(d)) |
| # The value expired. |
| self.mock_now(now, 2) |
| self.assertEqual(None, task_queues.probably_has_capacity(d)) |
| |
| def test_assert_bot_then_task(self): |
| self._assert_bot() |
| self._assert_task() |
| self.assertEqual([2980491642], task_queues.freshen_up_queues("bot1")) |
| |
| def test_assert_task_async_then_bot(self): |
| self._assert_task() |
| self._assert_bot() |
| self.assertEqual([2980491642], task_queues.freshen_up_queues("bot1")) |
| |
| def test_cleanup_after_bot(self): |
| self._assert_bot() |
| self._assert_task() |
| task_queues.cleanup_after_bot("bot1") |
| # BotInfo is deleted separately. |
| self.assert_count(0, task_queues.BotDimensionsMatches) |
| self.assert_count(1, bot_management.BotInfo) |
| |
| def test_assert_bot_dimensions_changed(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| self._assert_task( |
| { |
| "cpu": ["x86-64"], |
| "os": ["Amiga"], |
| "pool": ["default"], |
| } |
| ) |
| self._assert_bot(dimensions={"os": ["Amiga"]}) |
| |
| # The bot is matched to the task now. |
| self.assertEqual([2828582055], task_queues.freshen_up_queues("bot1")) |
| |
| # One hour later, the bot changes dimensions. |
| now += datetime.timedelta(hours=1) |
| self.mock_now(now) |
| self._assert_bot(dimensions={"os": "Commodore"}) |
| |
| # It is no longer matched to the task. |
| self.assertEqual([], task_queues.freshen_up_queues("bot1")) |
| |
| def test_assert_bot_queue_numbers(self): |
| self._assert_task( |
| { |
| "os": ["Ubuntu"], |
| "pool": ["default"], |
| } |
| ) |
| self._assert_task( |
| { |
| "os": ["Ubuntu"], |
| "id": ["bot-id"], |
| "pool": ["default"], |
| } |
| ) |
| |
| # This registers the bot for the first time and executes TQ tasks to match |
| # it to available tasks. Since TQ tasks happen after assert_bot call, there |
| # are no matching queues yet. |
| tq, queues = self._assert_bot("bot-id") |
| self.assertEqual(1, tq) |
| self.assertEqual([], queues) |
| |
| # Now all datastore structures are up-to-date and we can check what |
| # queues are actually matching the bot. |
| |
| # Matches both submitted tasks. |
| tq, queues = self._assert_bot("bot-id") |
| self.assertEqual(0, tq) |
| self.assertEqual([203088291, 3402762422], queues) |
| |
| def test_task_dimensions_expiry(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| self._assert_task() |
| self._assert_bot() |
| |
| # The bot is matched to the task now. |
| self.assertEqual([2980491642], task_queues.freshen_up_queues("bot1")) |
| |
| # Some time later the task dimensions set expires and gets cleaned up. |
| now += datetime.timedelta(hours=5) |
| self.mock_now(now) |
| task_queues.tidy_task_dimension_sets_async().get_result() |
| |
| # The bot is no longer matching the task. |
| self._assert_bot() |
| self.assertEqual([], task_queues.freshen_up_queues("bot1")) |
| |
| def test_hash_dimensions(self): |
| with self.assertRaises(AttributeError): |
| task_queues.hash_dimensions("this is not dict") |
| # Assert it doesn't return 0. |
| self.assertEqual(3649838548, task_queues.hash_dimensions({})) |
| |
| def test_hash_or_dimensions(self): |
| dim1 = ( |
| _gen_request( |
| properties=_gen_properties( |
| dimensions={"pool": ["pool-a"], "foo": ["a|c|b", "xyz"]} |
| ) |
| ) |
| .task_slice(0) |
| .properties.dimensions |
| ) |
| dim2 = ( |
| _gen_request( |
| properties=_gen_properties( |
| dimensions={"pool": ["pool-a"], "foo": ["xyz", "b|c|a"]} |
| ) |
| ) |
| .task_slice(0) |
| .properties.dimensions |
| ) |
| self.assertEqual( |
| task_queues.hash_dimensions(dim1), task_queues.hash_dimensions(dim2) |
| ) |
| |
| def test_expand_dimensions_to_flats(self): |
| expand = task_queues.expand_dimensions_to_flats |
| # Without OR |
| actual = set( |
| tuple(f) for f in expand({"foo": ["c", "a", "b"], "bar": ["x", "z", "y"]}) |
| ) |
| expected = { |
| ("bar:x", "bar:y", "bar:z", "foo:a", "foo:b", "foo:c"), |
| } |
| self.assertEqual(actual, expected) |
| # With OR |
| actual = set( |
| tuple(f) for f in expand({"foo": ["a|b|c", "def"], "bar": ["x|y", "z|w"]}) |
| ) |
| expected = { |
| ("bar:w", "bar:x", "foo:a", "foo:def"), |
| ("bar:w", "bar:x", "foo:b", "foo:def"), |
| ("bar:w", "bar:x", "foo:c", "foo:def"), |
| ("bar:w", "bar:y", "foo:a", "foo:def"), |
| ("bar:w", "bar:y", "foo:b", "foo:def"), |
| ("bar:w", "bar:y", "foo:c", "foo:def"), |
| ("bar:x", "bar:z", "foo:a", "foo:def"), |
| ("bar:x", "bar:z", "foo:b", "foo:def"), |
| ("bar:x", "bar:z", "foo:c", "foo:def"), |
| ("bar:y", "bar:z", "foo:a", "foo:def"), |
| ("bar:y", "bar:z", "foo:b", "foo:def"), |
| ("bar:y", "bar:z", "foo:c", "foo:def"), |
| } |
| # flats are sorted |
| actual = set( |
| tuple(f) |
| for f in expand( |
| { |
| "foo": ["a|y|b|z", "c"], |
| } |
| ) |
| ) |
| expected = { |
| ("foo:a", "foo:c"), |
| ("foo:b", "foo:c"), |
| ("foo:c", "foo:y"), |
| ("foo:c", "foo:z"), |
| } |
| self.assertEqual(actual, expected) |
| |
| def test_expiry_map(self): |
| sets = [ |
| {"dimensions": ("a:b", "a:c"), "expiry": 1663040000}, |
| {"dimensions": ("a:d", "a:e"), "expiry": 1663041111}, |
| ] |
| expiry_map = task_queues._sets_to_expiry_map(sets) |
| self.assertEqual(sets, task_queues._expiry_map_to_sets(expiry_map)) |
| without_ts = task_queues._expiry_map_to_sets(expiry_map, False) |
| self.assertEqual( |
| [ |
| {"dimensions": ("a:b", "a:c")}, |
| {"dimensions": ("a:d", "a:e")}, |
| ], |
| without_ts, |
| ) |
| |
| @staticmethod |
| def _stored_task_dims_expiry(expiry): |
| """Expiry of sets stored in TaskDimensionsSets given a task slice expiry.""" |
| expiry += datetime.timedelta(minutes=15) # mocked randomization |
| expiry += datetime.timedelta(hours=4) # extra time added in the txn |
| return expiry |
| |
| @parameterized.expand( |
| [ |
| ({"pool": ["p"], "k": ["v1"]}, "pool:p:1496061212"), |
| ({"pool": ["p"], "k": ["v1|v2"]}, "pool:p:3930364299"), |
| ({"pool": ["a:b"], "k": ["v1"]}, "pool:a%3Ab:2715169585"), |
| ({"id": ["b"], "pool": ["p"], "k": ["v1"]}, "bot:b:174937362"), |
| ({"id": ["b"], "pool": ["p"], "k": ["v1|v2"]}, "bot:b:903320086"), |
| ] |
| ) |
| def test_assert_task_dimensions_async(self, task_dims, expected_sets_id): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| expected_sets = task_queues.expand_dimensions_to_flats(task_dims) |
| |
| tq_tasks = [] |
| |
| @ndb.tasklet |
| def mocked_tq_task(task_sets_id, task_sets_dims, enqueued_ts): |
| self.assertEqual(task_sets_id, expected_sets_id) |
| self.assertEqual(task_sets_dims, expected_sets) |
| self.assertEqual(enqueued_ts, utils.utcnow()) |
| tq_tasks.append(1) |
| raise ndb.Return(True) |
| |
| self.mock(task_queues, "_tq_update_bot_matches_async", mocked_tq_task) |
| |
| def assert_task(expiry): |
| task_queues._assert_task_dimensions_async(task_dims, expiry).get_result() |
| self.execute_tasks() |
| seen = len(tq_tasks) |
| del tq_tasks[:] |
| return seen |
| |
| requested_exp = now + datetime.timedelta(hours=1) |
| stored_exp = self._stored_task_dims_expiry(requested_exp) |
| |
| # Add a never seen before task, should enqueue a TQ task. |
| self.assertEqual(assert_task(requested_exp), 1) |
| |
| # Check entities are correct and have correct expiry. |
| sets_key = ndb.Key(task_queues.TaskDimensionsSets, expected_sets_id) |
| sets_ent = sets_key.get() |
| self.assertEqual( |
| sets_ent.to_dict(), |
| { |
| "sets": [{"dimensions": dims} for dims in expected_sets], |
| }, |
| ) |
| info_key = ndb.Key(task_queues.TaskDimensionsInfo, 1, parent=sets_key) |
| info_ent = info_key.get() |
| self.assertEqual( |
| info_ent.to_dict(), |
| { |
| "sets": [ |
| {"dimensions": dims, "expiry": _to_timestamp(stored_exp)} |
| for dims in expected_sets |
| ], |
| "next_cleanup_ts": stored_exp + datetime.timedelta(minutes=5), |
| }, |
| ) |
| |
| # A bit later (before the stored expiry) add the exact same task. Should |
| # result in no TQ tasks and untouched entities. |
| now += datetime.timedelta(hours=1.5) |
| self.mock_now(now) |
| self.assertEqual(assert_task(now + datetime.timedelta(hours=1)), 0) |
| self.assertEqual(sets_key.get(), sets_ent) |
| self.assertEqual(info_key.get(), info_ent) |
| |
| # Adding a task even later results in the bump to the stored expiry (but no |
| # new TQ tasks). |
| now += datetime.timedelta(hours=6) |
| self.mock_now(now) |
| self.assertEqual(assert_task(now + datetime.timedelta(hours=1)), 0) |
| self.assertEqual( |
| info_key.get().next_cleanup_ts, |
| self._stored_task_dims_expiry(now + datetime.timedelta(hours=1)) |
| + datetime.timedelta(minutes=5), |
| ) |
| |
| def test_assert_task_dimensions_async_collision(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| # Simulate a hash collision, since it is hard (but not impossible) to get |
| # one for real. |
| self.mock(task_queues, "hash_dimensions", lambda _dims: 42) |
| |
| expected_sets_id = "pool:p:42" |
| sets_key = ndb.Key(task_queues.TaskDimensionsSets, expected_sets_id) |
| info_key = ndb.Key(task_queues.TaskDimensionsInfo, 1, parent=sets_key) |
| |
| tq_tasks = [] |
| |
| @ndb.tasklet |
| def mocked_tq_task(task_sets_id, task_sets_dims, enqueued_ts): |
| self.assertEqual(task_sets_id, expected_sets_id) |
| self.assertEqual(len(task_sets_dims), 1) |
| self.assertEqual(enqueued_ts, utils.utcnow()) |
| tq_tasks.append(tuple(task_sets_dims[0])) |
| raise ndb.Return(True) |
| |
| self.mock(task_queues, "_tq_update_bot_matches_async", mocked_tq_task) |
| |
| def assert_task(dim, exp): |
| task_queues._assert_task_dimensions_async( |
| { |
| "pool": ["p"], |
| "key": [dim], |
| }, |
| exp, |
| ).get_result() |
| self.execute_tasks() |
| seen = tq_tasks[:] |
| del tq_tasks[:] |
| return seen |
| |
| # Create a new TaskDimensionsSets. |
| exp1 = now + datetime.timedelta(hours=1) |
| stored_exp1 = self._stored_task_dims_expiry(exp1) |
| self.assertEqual(assert_task("1", exp1), [("key:1", "pool:p")]) |
| |
| # Append a new set to an existing TaskDimensionsSets. |
| now += datetime.timedelta(hours=1) |
| self.mock_now(now) |
| exp2 = now + datetime.timedelta(hours=1) |
| stored_exp2 = self._stored_task_dims_expiry(exp2) |
| self.assertEqual(assert_task("2", exp2), [("key:2", "pool:p")]) |
| |
| # Have both of them now. |
| self.assertEqual( |
| info_key.get().to_dict(), |
| { |
| "sets": [ |
| { |
| "dimensions": ["key:1", "pool:p"], |
| "expiry": _to_timestamp(stored_exp1), |
| }, |
| { |
| "dimensions": ["key:2", "pool:p"], |
| "expiry": _to_timestamp(stored_exp2), |
| }, |
| ], |
| "next_cleanup_ts": stored_exp1 + datetime.timedelta(minutes=5), |
| }, |
| ) |
| |
| # Refresh expiry of the second set, it should kick the first out as stale. |
| now = stored_exp2 - datetime.timedelta(hours=0.5) |
| self.mock_now(now) |
| exp3 = now + datetime.timedelta(hours=1) |
| stored_exp3 = self._stored_task_dims_expiry(exp3) |
| self.assertEqual(assert_task("2", exp3), []) |
| |
| # Have only one set now. |
| self.assertEqual( |
| info_key.get().to_dict(), |
| { |
| "sets": [ |
| { |
| "dimensions": ["key:2", "pool:p"], |
| "expiry": _to_timestamp(stored_exp3), |
| }, |
| ], |
| "next_cleanup_ts": stored_exp3 + datetime.timedelta(minutes=5), |
| }, |
| ) |
| |
| def test_update_bot_matches_async(self): |
| # Tested as a part of the overall workflow. |
| pass |
| |
| def test_rescan_matching_task_sets_async(self): |
| # Tested as a part of the overall workflow. |
| pass |
| |
| def test_tidy_task_dimension_sets_async(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| |
| # Use a transaction to be able to use _put_task_dimensions_sets_async that |
| # has an assert inside. We don't really care about atomicity here. |
| @ndb.transactional_tasklet(xg=True) |
| def prep(): |
| # Fully survives the cleanup. |
| yield task_queues._put_task_dimensions_sets_async( |
| "set0", |
| { |
| ("0:a",): now + datetime.timedelta(hours=2), |
| ("0:b",): now + datetime.timedelta(hours=2), |
| }, |
| ) |
| # Partially survives the cleanup. |
| yield task_queues._put_task_dimensions_sets_async( |
| "set1", |
| { |
| ("1:a",): now + datetime.timedelta(hours=1), |
| ("1:b",): now + datetime.timedelta(hours=2), |
| }, |
| ) |
| # Doesn't survive the cleanup. |
| yield task_queues._put_task_dimensions_sets_async( |
| "set2", |
| { |
| ("2:a",): now + datetime.timedelta(hours=1), |
| ("2:b",): now + datetime.timedelta(hours=1), |
| }, |
| ) |
| |
| prep().get_result() |
| |
| self.mock_now(now + datetime.timedelta(hours=1.5)) |
| self.assertTrue(task_queues.tidy_task_dimension_sets_async().get_result()) |
| |
| def fetch(sets_id): |
| sets_key = ndb.Key(task_queues.TaskDimensionsSets, sets_id) |
| info_key = ndb.Key(task_queues.TaskDimensionsInfo, 1, parent=sets_key) |
| sets, info = ndb.get_multi([sets_key, info_key]) |
| if not sets: |
| # Both must be missing at the same time. |
| self.assertIsNone(info) |
| return None |
| self.assertIsNotNone(info) |
| # Must contain same sets (sans `expiry`). |
| expless = [{"dimensions": s["dimensions"]} for s in info.sets] |
| self.assertEqual(sets.sets, expless) |
| return task_queues._sets_to_expiry_map(info.sets) |
| |
| # Fully survived. |
| self.assertEqual( |
| fetch("set0"), |
| { |
| ("0:a",): now + datetime.timedelta(hours=2), |
| ("0:b",): now + datetime.timedelta(hours=2), |
| }, |
| ) |
| # Partially survived. |
| self.assertEqual( |
| fetch("set1"), |
| { |
| ("1:b",): now + datetime.timedelta(hours=2), |
| }, |
| ) |
| # Fully gone. |
| self.assertIsNone(fetch("set2")) |
| |
| # Noop run for code coverage. |
| self.assertTrue(task_queues.tidy_task_dimension_sets_async().get_result()) |
| |
| # Full cleanup. |
| self.mock_now(now + datetime.timedelta(hours=2, minutes=6)) |
| self.assertTrue(task_queues.tidy_task_dimension_sets_async().get_result()) |
| |
| # All gone now. |
| self.assertIsNone(fetch("set0")) |
| self.assertIsNone(fetch("set1")) |
| |
| @staticmethod |
| def _create_task_dims_set(sets_id, *task_dims): |
| exp_map = { |
| tuple(dims): utils.utcnow() + datetime.timedelta(hours=10) |
| for dims in task_dims |
| } |
| ndb.transaction_async( |
| lambda: task_queues._put_task_dimensions_sets_async(sets_id, exp_map) |
| ).get_result() |
| |
| @staticmethod |
| def _delete_task_dims_set(sets_id): |
| ndb.transaction_async( |
| lambda: task_queues._delete_task_dimensions_sets_async(sets_id) |
| ).get_result() |
| |
| def test_assert_bot_dimensions_async(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| tq_tasks = [] |
| |
| @ndb.tasklet |
| def mocked_rescan(bot_id, rescan_counter, rescan_reason): |
| tq_tasks.append( |
| { |
| "bot_id": bot_id, |
| "rescan_counter": rescan_counter, |
| "rescan_reason": rescan_reason, |
| } |
| ) |
| raise ndb.Return(True) |
| |
| self.mock(task_queues, "_tq_rescan_matching_task_sets_async", mocked_rescan) |
| |
| def assert_no_tq_tasks(): |
| self.assertEqual(tq_tasks, []) |
| self.execute_tasks() |
| self.assertEqual(tq_tasks, []) |
| |
| def assert_tq_task(expected): |
| self.assertEqual(tq_tasks, []) |
| self.execute_tasks() |
| self.assertEqual(tq_tasks, [expected]) |
| del tq_tasks[:] |
| |
| dims = { |
| "id": ["bot-id"], |
| "pool": ["pool1", "pool2"], |
| "dim": ["0"], |
| } |
| bot_matches_key = ndb.Key(task_queues.BotDimensionsMatches, "bot-id") |
| |
| # The first call ever, no queues are assigned yet. |
| queues = task_queues._assert_bot_dimensions_async(dims).get_result() |
| self.assertEqual(queues, set()) |
| |
| # Created the entity. |
| bot_matches = bot_matches_key.get() |
| self.assertEqual( |
| bot_matches.to_dict(), |
| { |
| "dimensions": ["dim:0", "id:bot-id", "pool:pool1", "pool:pool2"], |
| "last_addition_ts": None, |
| "last_cleanup_ts": now, |
| "last_rescan_enqueued_ts": now, |
| "last_rescan_finished_ts": None, |
| "matches": [], |
| "next_rescan_ts": now + datetime.timedelta(minutes=30), |
| "rescan_counter": 1, |
| }, |
| ) |
| |
| # Enqueued the rescan task. |
| assert_tq_task( |
| { |
| "bot_id": "bot-id", |
| "rescan_counter": 1, |
| "rescan_reason": "dims added [dim:0 id:bot-id pool:pool1 pool:pool2]", |
| } |
| ) |
| |
| # Mock assignment of the queues. This happens in the TQ task usually. |
| self._create_task_dims_set("bot:bot-id:1", ["id:bot-id"]) |
| self._create_task_dims_set("pool:pool1:2", ["pool:pool1"]) |
| self._create_task_dims_set("pool:pool1:3", ["pool:pool1", "dim:0"]) |
| self._create_task_dims_set("pool:pool2:4", ["pool:pool2"]) |
| self._create_task_dims_set("pool:pool2:5", ["pool:pool2", "dim:0"]) |
| bot_matches.matches = [ |
| "bot:bot-id:1", |
| "pool:pool1:2", |
| "pool:pool1:3", |
| "pool:pool2:4", |
| "pool:pool2:5", |
| ] |
| bot_matches.put() |
| |
| # The next call discovers them and doesn't submit any TQ tasks. |
| queues = task_queues._assert_bot_dimensions_async(dims).get_result() |
| self.assertEqual( |
| queues, |
| { |
| "bot:bot-id:1", |
| "pool:pool1:2", |
| "pool:pool1:3", |
| "pool:pool2:4", |
| "pool:pool2:5", |
| }, |
| ) |
| assert_no_tq_tasks() |
| |
| # Some time later the rescan is triggered. |
| now += datetime.timedelta(minutes=31) |
| self.mock_now(now) |
| queues = task_queues._assert_bot_dimensions_async(dims).get_result() |
| self.assertEqual( |
| queues, |
| { |
| "bot:bot-id:1", |
| "pool:pool1:2", |
| "pool:pool1:3", |
| "pool:pool2:4", |
| "pool:pool2:5", |
| }, |
| ) |
| |
| # Enqueued the rescan task. |
| assert_tq_task( |
| { |
| "bot_id": "bot-id", |
| "rescan_counter": 2, |
| "rescan_reason": "periodic", |
| } |
| ) |
| |
| # Bots dimensions change. Unmatching queues are no longer reported. |
| dims["dim"] = ["1"] |
| queues = task_queues._assert_bot_dimensions_async(dims).get_result() |
| self.assertEqual( |
| queues, |
| { |
| "bot:bot-id:1", |
| "pool:pool1:2", |
| "pool:pool2:4", |
| }, |
| ) |
| |
| # Enqueued the rescan task. |
| assert_tq_task( |
| { |
| "bot_id": "bot-id", |
| "rescan_counter": 3, |
| "rescan_reason": "dims added [dim:1], dims removed [dim:0]", |
| } |
| ) |
| |
| # Changes are reflected in the entity. |
| bot_matches = bot_matches_key.get() |
| self.assertEqual( |
| bot_matches.to_dict(), |
| { |
| "dimensions": ["dim:1", "id:bot-id", "pool:pool1", "pool:pool2"], |
| "last_addition_ts": None, |
| "last_cleanup_ts": now, |
| "last_rescan_enqueued_ts": now, |
| "last_rescan_finished_ts": None, # we mocked this out |
| "matches": ["bot:bot-id:1", "pool:pool1:2", "pool:pool2:4"], |
| "next_rescan_ts": now + datetime.timedelta(minutes=30), |
| "rescan_counter": 3, |
| }, |
| ) |
| |
| # Some task set disappears (e.g. gets cleaned up by the cron). |
| now += datetime.timedelta(minutes=5) |
| self.mock_now(now) |
| self._delete_task_dims_set("pool:pool2:4") |
| |
| # It is no longer assigned to the bot. This doesn't enqueue a task. |
| queues = task_queues._assert_bot_dimensions_async(dims).get_result() |
| self.assertEqual( |
| queues, |
| { |
| "bot:bot-id:1", |
| "pool:pool1:2", |
| }, |
| ) |
| assert_no_tq_tasks() |
| |
| # Changes are reflected in the entity. |
| prev_state = bot_matches.to_dict() |
| bot_matches = bot_matches_key.get() |
| self.assertEqual( |
| bot_matches.to_dict(), |
| { |
| "dimensions": ["dim:1", "id:bot-id", "pool:pool1", "pool:pool2"], |
| "last_addition_ts": None, |
| "last_cleanup_ts": now, |
| "last_rescan_enqueued_ts": prev_state["last_rescan_enqueued_ts"], |
| "last_rescan_finished_ts": None, # we mocked this out |
| "matches": ["bot:bot-id:1", "pool:pool1:2"], |
| "next_rescan_ts": prev_state["next_rescan_ts"], |
| "rescan_counter": prev_state["rescan_counter"], |
| }, |
| ) |
| |
| def test_tq_rescan_matching_task_sets_async(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| # Prepare all active task dims sets. |
| self._create_task_dims_set("bot:bot-id:1", ["id:bot-id"]) |
| self._create_task_dims_set("bot:bot-id:2", ["id:bot-id", "dim:0"]) |
| self._create_task_dims_set("bot:bot-id:3", ["id:bot-id", "dim:1"]) |
| self._create_task_dims_set("pool:pool1:4", ["pool:pool1"]) |
| self._create_task_dims_set("pool:pool1:5", ["pool:pool1", "dim:0"]) |
| self._create_task_dims_set("pool:pool1:6", ["pool:pool1", "dim:1"]) |
| self._create_task_dims_set("pool:pool2:7", ["pool:pool2"]) |
| self._create_task_dims_set("pool:pool2:8", ["pool:pool2", "dim:0"]) |
| self._create_task_dims_set("pool:pool2:9", ["pool:pool2", "dim:1"]) |
| |
| # Put some extra garbage that must no be visited since it is outside of the |
| # range of datastore scan query. |
| self._create_task_dims_set("bot:bot-id:", ["id:bot-id"]) |
| self._create_task_dims_set("bot:bot-id:/", ["id:bot-id"]) |
| self._create_task_dims_set("bot:bot-id::", ["id:bot-id"]) |
| self._create_task_dims_set("pool:pool1:", ["pool:pool1"]) |
| self._create_task_dims_set("pool:pool1:/", ["pool:pool1"]) |
| self._create_task_dims_set("pool:pool1::", ["pool:pool1"]) |
| |
| # Prepare BotDimensionsMatches in some initial pre-scan state: it has new |
| # dimensions (dim:1), but matches are still for old ones (dim:0). |
| task_queues.BotDimensionsMatches( |
| id="bot-id", |
| dimensions=["dim:1", "id:bot-id", "pool:pool1", "pool:pool2"], |
| matches=[ |
| "bot:bot-id:1", |
| "bot:bot-id:2", |
| "bot:bot-id:444", # missing, should be unmatched |
| "pool:missing:555", # missing, should be unmatched |
| "pool:pool1:4", |
| "pool:pool1:5", |
| "pool:pool2:7", |
| "pool:pool2:8", |
| ], |
| last_rescan_enqueued_ts=now, |
| next_rescan_ts=now + datetime.timedelta(hours=1), |
| rescan_counter=555, |
| ).put() |
| |
| # Run the rescan. |
| self.assertTrue( |
| task_queues._tq_rescan_matching_task_sets_async( |
| "bot-id", 555, "reason" |
| ).get_result() |
| ) |
| |
| # The state was updated correctly to match on `dim:1`. Old entries were |
| # kicked out. |
| ent = ndb.Key(task_queues.BotDimensionsMatches, "bot-id").get() |
| self.assertEqual( |
| ent.to_dict(), |
| { |
| "dimensions": ["dim:1", "id:bot-id", "pool:pool1", "pool:pool2"], |
| "last_addition_ts": None, |
| "last_cleanup_ts": now, |
| "last_rescan_enqueued_ts": now, |
| "last_rescan_finished_ts": now, |
| "matches": [ |
| "bot:bot-id:1", |
| "bot:bot-id:3", |
| "pool:pool1:4", |
| "pool:pool1:6", |
| "pool:pool2:7", |
| "pool:pool2:9", |
| ], |
| "next_rescan_ts": now + datetime.timedelta(hours=1), |
| "rescan_counter": 555, |
| }, |
| ) |
| |
| def test_tq_update_bot_matches_async_pool(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| task_sets_id = "pool:p:1" |
| task_sets_dims = [["pool:p", "dim:0"], ["pool:p", "dim:1"]] |
| self._create_task_dims_set(task_sets_id, *task_sets_dims) |
| |
| def register_bot(bot_id, dims, matches): |
| task_queues.BotDimensionsMatches( |
| id=bot_id, |
| dimensions=sorted(dims + ["id:" + bot_id]), |
| matches=matches, |
| last_addition_ts=utils.EPOCH, |
| ).put() |
| |
| # Bots that should be matched. |
| register_bot("match-0", ["pool:p", "dim:0"], ["pool:p:555"]) |
| register_bot("match-1", ["pool:p", "dim:1"], []) |
| register_bot("match-2", ["pool:p", "dim:0", "dim:1"], []) |
| register_bot("match-3", ["pool:p", "dim:0", "extra:0"], []) |
| |
| # Will be visited, but not updated (since it already has the match) |
| register_bot("untouched-0", ["pool:p", "dim:0"], [task_sets_id]) |
| |
| # Bots that should be left untouched. |
| register_bot("mismatch-0", ["pool:p"], ["pool:p:555"]) |
| register_bot("mismatch-1", ["pool:p", "dim:nope"], ["pool:p:555"]) |
| register_bot("mismatch-2", ["pool:another", "dim:0"], ["pool:p:555"]) |
| |
| self.assertTrue( |
| task_queues._tq_update_bot_matches_async( |
| task_sets_id, task_sets_dims, now |
| ).get_result() |
| ) |
| |
| def assert_matched(bot_id, matches): |
| ent = ndb.Key(task_queues.BotDimensionsMatches, bot_id).get() |
| self.assertEqual(ent.matches, matches) |
| self.assertEqual(ent.last_addition_ts, now) |
| |
| def assert_untouched(bot_id, matches): |
| ent = ndb.Key(task_queues.BotDimensionsMatches, bot_id).get() |
| self.assertEqual(ent.matches, matches) |
| self.assertEqual(ent.last_addition_ts, utils.EPOCH) |
| |
| assert_matched("match-0", [task_sets_id, "pool:p:555"]) |
| assert_matched("match-1", [task_sets_id]) |
| assert_matched("match-2", [task_sets_id]) |
| assert_matched("match-3", [task_sets_id]) |
| |
| assert_untouched("untouched-0", [task_sets_id]) |
| assert_untouched("mismatch-0", ["pool:p:555"]) |
| assert_untouched("mismatch-1", ["pool:p:555"]) |
| assert_untouched("mismatch-2", ["pool:p:555"]) |
| |
| def test_tq_update_bot_matches_async_bot(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| task_queues.BotDimensionsMatches( |
| id="bot-id", |
| dimensions=["id:bot-id", "dim:0", "dim:1", "pool:p"], |
| matches=["pool:p:123", "bot:bot-id:456"], |
| last_addition_ts=utils.EPOCH, |
| ).put() |
| |
| def match_task(task_sets_id, task_sets_dims): |
| self._create_task_dims_set(task_sets_id, *task_sets_dims) |
| self.assertTrue( |
| task_queues._tq_update_bot_matches_async( |
| task_sets_id, task_sets_dims, now |
| ).get_result() |
| ) |
| |
| # Will be matched. |
| match_task("bot:bot-id:1", [["id:bot-id", "dim:0"]]) |
| match_task("bot:bot-id:2", [["id:bot-id", "dim:1"]]) |
| match_task("bot:bot-id:3", [["id:bot-id", "dim:0", "dim:1"]]) |
| match_task("bot:bot-id:4", [["id:bot-id", "dim:z"], ["id:bot-id", "dim:0"]]) |
| match_task("bot:bot-id:5", [["id:bot-id"]]) |
| |
| # Will not be matched, wrong dims. |
| match_task("bot:bot-id:6", [["id:bot-id", "dim:3"]]) |
| match_task("bot:bot-id:7", [["id:bot-id", "dim:0", "extra:1"]]) |
| |
| ent = ndb.Key(task_queues.BotDimensionsMatches, "bot-id").get() |
| self.assertEqual( |
| ent.matches, |
| [ |
| "bot:bot-id:1", |
| "bot:bot-id:2", |
| "bot:bot-id:3", |
| "bot:bot-id:4", |
| "bot:bot-id:456", |
| "bot:bot-id:5", |
| "pool:p:123", |
| ], |
| ) |
| |
| def test_tq_update_bot_matches_async_skip(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| task_queues.BotDimensionsMatches( |
| id="bot-id", |
| dimensions=["id:bot-id", "dim:0", "dim:1", "pool:p"], |
| matches=[], |
| last_addition_ts=utils.EPOCH, |
| ).put() |
| |
| def matches(): |
| return ndb.Key(task_queues.BotDimensionsMatches, "bot-id").get().matches |
| |
| def call(sets_id, tasks_dims): |
| return task_queues._tq_update_bot_matches_async( |
| sets_id, tasks_dims, now |
| ).get_result() |
| |
| # Will do nothing, no TaskDimensionsSets entity is stored. |
| self.assertTrue(call("bot:bot-id:1", [["id:bot-id"]])) |
| self.assertEqual(matches(), []) |
| |
| # Will do nothing, all sets are not in the entity. |
| self._create_task_dims_set("bot:bot-id:1", ["id:bot-id", "dim:0"]) |
| self.assertTrue( |
| call("bot:bot-id:1", [["id:bot-id"], ["id:bot-id", "dim:1"]]) |
| ) |
| self.assertEqual(matches(), []) |
| |
| def test_tq_update_bot_matches_async_no_matches(self): |
| now = datetime.datetime(2010, 1, 2, 3, 4, 5) |
| self.mock_now(now) |
| |
| def call(sets_id, tasks_dims): |
| return task_queues._tq_update_bot_matches_async( |
| sets_id, tasks_dims, now |
| ).get_result() |
| |
| # No matching bot. |
| self._create_task_dims_set("bot:bot-id:1", ["id:bot-id"]) |
| self.assertTrue(call("bot:bot-id:1", [["id:bot-id"]])) |
| |
| # No matches in a pool. |
| self._create_task_dims_set("pool:p:1", ["pool:p"]) |
| self.assertTrue(call("pool:p:1", [["pool:p"]])) |
| |
| def test_tidy_bot_dimensions_matches_async(self): |
| for i in range(0, 10): |
| self._assert_bot("bot-%d" % i) |
| |
| @ndb.tasklet |
| def should_keep_async(bot_id): |
| idx = int(bot_id[len("bot-") :]) |
| raise ndb.Return(idx % 2 == 0) |
| |
| ok = task_queues.tidy_bot_dimensions_matches_async( |
| should_keep_async |
| ).get_result() |
| self.assertTrue(ok) |
| |
| remains = list(task_queues.BotDimensionsMatches.query()) |
| self.assertEqual( |
| ["bot-0", "bot-2", "bot-4", "bot-6", "bot-8"], |
| [ent.key.string_id() for ent in remains], |
| ) |
| |
| |
| class TestMapAsync(test_env_handlers.AppTestBase): |
| # Page size in queries. |
| PAGE = 4 |
| |
| @staticmethod |
| def populate(bot_dims): |
| ndb.put_multi( |
| [ |
| task_queues.BotDimensionsMatches(id=bot_id, dimensions=dims) |
| for bot_id, dims in bot_dims.items() |
| ] |
| ) |
| return sorted(bot_dims.keys()) |
| |
| @staticmethod |
| def query(dim=None): |
| q = task_queues.BotDimensionsMatches.query() |
| if dim: |
| q = q.filter(task_queues.BotDimensionsMatches.dimensions == dim) |
| return q, task_queues._Logger("") |
| |
| @staticmethod |
| def call(queries, cb, max_pages=None, keys_only=False): |
| return task_queues._map_async( |
| queries, |
| cb, |
| max_concurrency=3, |
| page_size=TestMapAsync.PAGE, |
| max_pages=max_pages, |
| keys_only=keys_only, |
| ).get_result() |
| |
| def test_single_query_sync_cb(self): |
| seen = [] |
| |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| |
| bots = self.populate({"bot%d" % i: [] for i in range(10)}) |
| self.assertTrue(self.call([self.query()], cb)) |
| self.assertEqual(sorted(seen), bots) |
| |
| def test_single_query_sync_cb_keys_only(self): |
| seen = [] |
| |
| def cb(key): |
| seen.append(key.string_id()) |
| |
| bots = self.populate({"bot%d" % i: [] for i in range(10)}) |
| self.assertTrue(self.call([self.query()], cb, keys_only=True)) |
| self.assertEqual(sorted(seen), bots) |
| |
| def test_many_queries_sync_cb(self): |
| seen = [] |
| |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| |
| bots = [] |
| bots.extend(self.populate({"bot1-%d" % i: ["k:v1"] for i in range(10)})) |
| bots.extend(self.populate({"bot2-%d" % i: ["k:v2"] for i in range(10)})) |
| bots.extend(self.populate({"bot3-%d" % i: ["k:v3"] for i in range(10)})) |
| |
| queries = [self.query("k:v1"), self.query("k:v2"), self.query("k:v3")] |
| self.assertTrue(self.call(queries, cb)) |
| self.assertEqual(sorted(seen), bots) |
| |
| def test_single_query_async_cb(self): |
| seen = [] |
| |
| @ndb.tasklet |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| yield ndb.sleep(random.uniform(0.0, 0.1)) |
| raise ndb.Return(True) |
| |
| bots = self.populate({"bot%d" % i: [] for i in range(10)}) |
| self.assertTrue(self.call([self.query()], cb)) |
| self.assertEqual(sorted(seen), bots) |
| |
| def test_many_queries_async_cb(self): |
| seen = [] |
| |
| @ndb.tasklet |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| yield ndb.sleep(random.uniform(0.0, 0.1)) |
| raise ndb.Return(True) |
| |
| bots = [] |
| bots.extend(self.populate({"bot1-%d" % i: ["k:v1"] for i in range(10)})) |
| bots.extend(self.populate({"bot2-%d" % i: ["k:v2"] for i in range(10)})) |
| bots.extend(self.populate({"bot3-%d" % i: ["k:v3"] for i in range(10)})) |
| |
| queries = [self.query("k:v1"), self.query("k:v2"), self.query("k:v3")] |
| self.assertTrue(self.call(queries, cb)) |
| self.assertEqual(sorted(seen), bots) |
| |
| def test_query_timeout(self): |
| seen = [] |
| |
| @ndb.tasklet |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| yield ndb.sleep(random.uniform(0.0, 0.1)) |
| raise ndb.Return(True) |
| |
| b1 = self.populate({"b1-%d" % i: ["k:v1"] for i in range(self.PAGE * 2)}) |
| b2 = self.populate({"b2-%d" % i: ["k:v2"] for i in range(self.PAGE * 2)}) |
| b3 = self.populate({"b3-%d" % i: ["k:v3"] for i in range(self.PAGE * 3)}) |
| |
| queries = [self.query("k:v1"), self.query("k:v2"), self.query("k:v3")] |
| self.assertFalse(self.call(queries, cb, max_pages=2)) |
| |
| # The last page if k:v3 query is missing. |
| expected = [] |
| expected.extend(b1) |
| expected.extend(b2) |
| expected.extend(b3[: self.PAGE * 2]) |
| self.assertEqual(sorted(seen), expected) |
| |
| def test_processing_error(self): |
| seen = [] |
| |
| @ndb.tasklet |
| def cb(ent): |
| seen.append(ent.key.string_id()) |
| yield ndb.sleep(random.uniform(0.0, 0.1)) |
| raise ndb.Return(ent.key.string_id() != "bot5") |
| |
| bots = self.populate({"bot%d" % i: [] for i in range(10)}) |
| self.assertFalse(self.call([self.query()], cb)) |
| self.assertEqual(sorted(seen), bots) |
| |
| |
| 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() |
