appengine/swarming/server/named_caches_test.py

#!/usr/bin/env vpython
# Copyright 2018 The LUCI Authors. All rights reserved.
# Use of this source code is governed under the Apache License, Version 2.0
# that can be found in the LICENSE file.

import datetime
import hashlib
import json
import logging
import sys
import unittest

# pylint: disable=wrong-import-position
import test_env

test_env.setup_test_env()

from google.appengine.ext import ndb

from components import utils
from server import bot_management
from server import named_caches
from server import pools_config
from test_support import test_case


def _bot_event(bot_id, pool, caches, oses):
  """Calls bot_management.bot_event with default arguments."""
  dimensions = {
    "id": [bot_id],
    "os": oses or ["Linux", "Ubuntu", "Ubuntu-16.04"],
    "pool": [pool],
  }
  # Format is named_caches: {name: [['shortname', size], timestamp]}.
  state = {
    "named_caches": {name: [["a", size], 10] for name, size in caches.items()}
  }
  bot_management.bot_event(
    event_type="request_sleep",
    bot_id=bot_id,
    external_ip="8.8.4.4",
    authenticated_as="bot:%s.domain" % bot_id,
    dimensions=dimensions,
    state=state or {"ram": 65},
    version=unicode(hashlib.sha256().hexdigest()),
    register_dimensions=True,
  )


class NamedCachesTest(test_case.TestCase):
  APP_DIR = test_env.APP_DIR

  def setUp(self):
    super(NamedCachesTest, self).setUp()
    self.mock(utils, "enqueue_task", self._enqueue_task)
    self.mock(pools_config, "known", lambda: ["first", "second"])

  @ndb.non_transactional
  def _enqueue_task(self, url, queue_name, payload):
    if queue_name == "named-cache-task":
      params = json.loads(payload)
      self.assertEqual(True, named_caches.task_update_pool(params["pool"]))
      return True
    self.fail(url)
    return False

  def test_simple(self):
    _bot_event("first1", "first", {"git": 1}, None)
    _bot_event("first2", "first", {"build": 100000, "git": 1000}, None)
    self.assertEqual(2, named_caches.cron_update_named_caches())

    self.assertEqual(2, named_caches.NamedCache.query().count())
    oses = ["Linux"]
    hints = named_caches.get_hints("first", oses, ["git", "build", "new"])
    self.assertEqual([1000, 100000, -1], hints)

  def test_p95(self):
    # Create 45 bots with cache 'foo' size between 1 and 45.
    for i in range(45):
      _bot_event("second%d" % i, "second", {"foo": i + 1}, None)
    self.assertEqual(2, named_caches.cron_update_named_caches())

    self.assertEqual(1, named_caches.NamedCache.query().count())
    oses = ["Linux"]
    hints = named_caches.get_hints("second", oses, ["foo"])
    # Roughly p95.
    self.assertEqual([43], hints)

  def test_fuzzy_other_os(self):
    # Use the hint from 'Mac' (the larger one) even if requesting for Linux.
    _bot_event("first1", "first", {"build": 50000}, ["Android"])
    _bot_event("first2", "first", {"build": 100000}, ["Mac"])
    self.assertEqual(2, named_caches.cron_update_named_caches())

    self.assertEqual(2, named_caches.NamedCache.query().count())
    oses = ["Linux"]
    hints = named_caches.get_hints("first", oses, ["build"])
    self.assertEqual([100000], hints)

  def test_expired(self):
    now = datetime.datetime(2015, 1, 1, 1, 1, 1)
    self.mock_now(now)
    _bot_event("first1", "first", {"build": 100000}, ["Mac"])
    self.assertEqual(2, named_caches.cron_update_named_caches())
    self.assertEqual([now], [e.ts for e in named_caches.NamedCache.query()])

    self.mock_now(now, 8 * 24 * 60 * 60)
    self.assertEqual(2, named_caches.cron_update_named_caches())
    self.assertEqual(1, named_caches.NamedCache.query().count())
    self.assertEqual([now], [e.ts for e in named_caches.NamedCache.query()])

    self.mock_now(now, 8 * 24 * 60 * 60 + 1)
    self.assertEqual(2, named_caches.cron_update_named_caches())
    self.assertEqual(0, named_caches.NamedCache.query().count())


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()