Expose named caches as dimensions.

appengine/swarming/local_smoke_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.
 
 """Integration test for the Swarming server, Swarming bot and Swarming client.
 
 It starts both a Swarming server and a Swarming bot and triggers tasks with the
 Swarming client to ensure the system works end to end.
 """
 
 import base64
 import json
 import logging
 import os
 import re
 import signal
 import socket
 import sys
 import tempfile
 import textwrap
+import time
 import unittest
 import urllib
 
 APP_DIR = os.path.dirname(os.path.abspath(__file__))
 BOT_DIR = os.path.join(APP_DIR, 'swarming_bot')
 CLIENT_DIR = os.path.join(APP_DIR, '..', '..', 'client')
 
 from tools import start_bot
 from tools import start_servers
 
@@ skipping 169 matching lines @@
       os.remove(tmp)
 
   def terminate(self, bot_id):
     return self._run('terminate', ['--wait', bot_id])
 
   def cleanup(self):
     if self._tmpdir:
       file_path.rmtree(self._tmpdir)
       self._tmpdir = None
 
+  def query_bot(self):
+    """Returns the bot's properties."""
+    data = json.loads(self._capture('query', ['bots/list', '--limit', '10']))
+    if not data.get('items'):
+      return None
+    assert len(data['items']) == 1
+    return data['items'][0]
+
   def dump_log(self):
     print >> sys.stderr, '-' * 60
     print >> sys.stderr, 'Client calls'
     print >> sys.stderr, '-' * 60
     for i in xrange(self._index):
       with open(os.path.join(self._tmpdir, 'client_%d.log' % i), 'rb') as f:
         log = f.read().strip('\n')
       for l in log.splitlines():
         sys.stderr.write('  %s\n' % l)
 
@@ skipping 13 matching lines @@
       '--verbose',
     ] + args
     with open(name, 'wb') as f:
       f.write('\nRunning: %s\n' % ' '.join(cmd))
       f.flush()
       p = subprocess42.Popen(
           cmd, stdout=f, stderr=subprocess42.STDOUT, cwd=CLIENT_DIR)
       p.communicate()
       return p.returncode
 
+  def _capture(self, command, args):
+    cmd = [
+      sys.executable, 'swarming.py', command, '-S', self._swarming_server,
+    ] + args
+    p = subprocess42.Popen(cmd, stdout=subprocess42.PIPE, cwd=CLIENT_DIR)
+    return p.communicate()[0]
+
 
 def gen_expected(**kwargs):
   expected = {
     u'abandoned_ts': None,
     u'bot_dimensions': None,
     u'bot_id': unicode(socket.getfqdn().split('.', 1)[0]),
     u'children_task_ids': [],
     u'cost_saved_usd': None,
     u'deduped_from': None,
     u'exit_codes': [0],
@@ skipping 19 matching lines @@
   # convert_to_old_format.py in //client/swarming.py.
   keys = set(expected) | {u'performance_stats'}
   assert keys.issuperset(kwargs)
   expected.update({unicode(k): v for k, v in kwargs.iteritems()})
   return expected
 
 
 class Test(unittest.TestCase):
   maxDiff = None
   client = None
-  dimensions = None
   servers = None
   bot = None
 
-  @classmethod
-  def setUpClass(cls):
-    cls.dimensions = os_utilities.get_dimensions()
-
   def setUp(self):
     super(Test, self).setUp()
-    # Reset the bot's cache at the start of each task, so that the cache reuse
-    # data becomes deterministic.
-    # Main caveat is 'isolated_upload' as the isolate server is not cleared.
-    self.bot.wipe_cache()
+    self.dimensions = os_utilities.get_dimensions()
+    # Reset the bot's isolated cache at the start of each task, so that the
+    # cache reuse data becomes deterministic. Only restart the bot when it had a
+    # named cache because it takes multiple seconds to to restart the bot. :(
+    #
+    # TODO(maruel): 'isolated_upload' is not deterministic because the isolate
+    # server not cleared.
+    old = self.client.query_bot()
+    started_ts = json.loads(old['state'])['started_ts'] if old else None
+    logging.info('setUp: started_ts was %s', started_ts)
+    had_cache = any(
+        u'caches' == i['key'] for i in old['dimensions']) if old else False
+    self.bot.wipe_cache(had_cache)
+    # The bot restarts due to wipe_cache() so wait for the bot to come back
+    # online. It may takes a few loop.
+    while True:
+      state = self.client.query_bot()
+      if not state:
+        time.sleep(0.1)
+        continue
+      if not had_cache:
+        break
+      new_started_ts = json.loads(state['state'])['started_ts']
+      logging.info('setUp: new_started_ts is %s', new_started_ts)
+      # This assumes that starting the bot and running the previous test case
+      # took more than 1s.
+      if not started_ts or new_started_ts != started_ts:
+        dimensions = {i['key']: i['value'] for i in state['dimensions']}
+        self.assertNotIn(u'caches', dimensions)
+        break
 
   def gen_expected(self, **kwargs):
     return gen_expected(bot_dimensions=self.dimensions, **kwargs)
 
   def test_raw_bytes(self):
     # A string of a letter 'A', UTF-8 BOM then UTF-16 BOM then UTF-EDBCDIC then
     # invalid UTF-8 and the letter 'B'. It is double escaped so it can be passed
     # down the shell.
     invalid_bytes = 'A\\xEF\\xBB\\xBF\\xFE\\xFF\\xDD\\x73\\x66\\x73\\xc3\\x28B'
     args = [
@@ skipping 322 matching lines @@
       self._run_isolated(
           hello_world, 'secret_bytes',
           ['--secret-bytes-path', tmp, '--', '${ISOLATED_OUTDIR}'],
           expected_summary, {os.path.join('0', 'sekret'): 'foobar\n'})
     finally:
       os.remove(tmp)
 
   def test_local_cache(self):
     # First task creates the cache, second copy the content to the output
     # directory. Each time it's the exact same script.
+    dimensions = {
+        i['key']: i['value'] for i in self.client.query_bot()['dimensions']}
+    self.assertEqual(set(self.dimensions), set(dimensions))
+    self.assertNotIn(u'cache', set(dimensions))
     script = '\n'.join((
       'import os, shutil, sys',
       'p = "p/b/a.txt"',
       'if not os.path.isfile(p):',
       '  with open(p, "wb") as f:',
       '    f.write("Yo!")',
       'else:',
       '  shutil.copy(p, sys.argv[1])',
       'print "hi"'))
     sizes = sorted([len(script), 200])
@@ skipping 36 matching lines @@
           u'initial_size': unicode(sum(sizes)),
           u'items_cold': [],
           u'items_hot': sizes,
         },
         u'isolated_upload': {
           u'items_cold': [3, 110],
           u'items_hot': [],
         },
       },
     )
+    # The previous task caused the bot to have a named cache.
+    expected_summary['bot_dimensions'] = (
+        expected_summary['bot_dimensions'].copy())
+    expected_summary['bot_dimensions'][u'caches'] = [u'fuu']
     self._run_isolated(
         script, 'cache_second',
         ['--named-cache', 'fuu', 'p/b', '--', '${ISOLATED_OUTDIR}/yo'],
         expected_summary,
         {'0/yo': 'Yo!'})
 
+    # Check that the bot now has a cache dimension by independently querying.
+    expected = set(self.dimensions)
+    expected.add(u'caches')
+    dimensions = {
+        i['key']: i['value'] for i in self.client.query_bot()['dimensions']}
+    self.assertEqual(expected, set(dimensions))
+
   def _run_isolated(self, hello_world, name, args, expected_summary,
       expected_files, deduped=False, isolated_content=None):
     """Runs hello_world.py as an isolated file."""
     # Shared code for all test_isolated_* test cases.
     tmpdir = tempfile.mkdtemp(prefix='swarming_smoke')
     try:
       isolate_path = os.path.join(tmpdir, 'i.isolate')
       isolated_path = os.path.join(tmpdir, 'i.isolated')
       with open(isolate_path, 'wb') as f:
         json.dump(isolated_content or ISOLATE_HELLO_WORLD, f)
@@ skipping 99 matching lines @@
   finally:
     if client and not leak:
       client.cleanup()
 
 
 def main():
   fix_encoding.fix_encoding()
   verbose = '-v' in sys.argv
   leak = bool('--leak' in sys.argv)
   if leak:
+    # Note that --leak will not guarantee that 'c' and 'isolated_cache' are
+    # kept. Only the last test case will leak these two directories.
     sys.argv.remove('--leak')
   if verbose:
     logging.basicConfig(level=logging.INFO)
     Test.maxDiff = None
   else:
     logging.basicConfig(level=logging.ERROR)
 
   # Force language to be English, otherwise the error messages differ from
   # expectations.
   os.environ['LANG'] = 'en_US.UTF-8'
@@ skipping 40 matching lines @@
     if bot is not None and bot.poll() is None:
       bot.kill()
       bot.wait()
   finally:
     cleanup(bot, client, servers, failed or verbose, leak)
   return int(failed)
 
 
 if __name__ == '__main__':
   sys.exit(main())