src/tools/swarming_client/tests/isolateserver_smoke_test.py

   1 #!/usr/bin/env python
   2 # Copyright 2013 The Swarming Authors. All rights reserved.
   3 # Use of this source code is governed under the Apache License, Version 2.0 that
   4 # can be found in the LICENSE file.
   5
   6 import hashlib
   7 import logging
   8 import os
   9 import shutil
  10 import subprocess
  11 import sys
  12 import tempfile
  13 import time
  14 import unittest
  15
  16 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  17 sys.path.insert(0, ROOT_DIR)
  18
  19 import isolated_format
  20
  21 ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  22
  23 # Ensure that the testing machine has access to this server.
  24 ISOLATE_SERVER = 'https://isolateserver.appspot.com/'
  25
  26 # The directory containing the test data files.
  27 TEST_DATA_DIR = os.path.join(ROOT_DIR, 'tests', 'isolateserver')
  28
  29
  30 class IsolateServerArchiveSmokeTest(unittest.TestCase):
  31   def setUp(self):
  32     super(IsolateServerArchiveSmokeTest, self).setUp()
  33     # The namespace must end in '-gzip' since all files are now compressed
  34     # before being uploaded.
  35     # TODO(maruel): This should not be leaked to the client. It's a
  36     # transport/storage detail.
  37     self.namespace = ('temporary' + str(long(time.time())).split('.', 1)[0]
  38                       + '-gzip')
  39     self.rootdir = tempfile.mkdtemp(prefix='isolateserver')
  40
  41   def tearDown(self):
  42     try:
  43       shutil.rmtree(self.rootdir)
  44     finally:
  45       super(IsolateServerArchiveSmokeTest, self).tearDown()
  46
  47   def _run(self, args):
  48     """Runs isolateserver.py."""
  49     cmd = [
  50         sys.executable,
  51         os.path.join(ROOT_DIR, 'isolateserver.py'),
  52     ]
  53     cmd.extend(args)
  54     cmd.extend(
  55         [
  56           '--isolate-server', ISOLATE_SERVER,
  57           '--namespace', self.namespace
  58         ])
  59     if '-v' in sys.argv:
  60       cmd.append('--verbose')
  61       subprocess.check_call(cmd)
  62     else:
  63       subprocess.check_output(cmd)
  64
  65   def _archive_given_files(self, files):
  66     """Given a list of files, call isolateserver.py with them. Then
  67     verify they are all on the server."""
  68     files = [os.path.join(TEST_DATA_DIR, filename) for filename in files]
  69     self._run(['archive'] + files)
  70     self._download_given_files(files)
  71
  72   def _download_given_files(self, files):
  73     """Tries to download the files from the server."""
  74     args = ['download', '--target', self.rootdir]
  75     file_hashes = [isolated_format.hash_file(f, hashlib.sha1) for f in files]
  76     for f in file_hashes:
  77       args.extend(['--file', f, f])
  78     self._run(args)
  79     # Assert the files are present.
  80     actual = [
  81         isolated_format.hash_file(os.path.join(self.rootdir, f), hashlib.sha1)
  82         for f in os.listdir(self.rootdir)
  83     ]
  84     self.assertEqual(sorted(file_hashes), sorted(actual))
  85
  86   def test_archive_empty_file(self):
  87     self._archive_given_files(['empty_file.txt'])
  88
  89   def test_archive_small_file(self):
  90     self._archive_given_files(['small_file.txt'])
  91
  92   def test_archive_huge_file(self):
  93     # Create a file over 2gbs.
  94     filepath = None
  95     try:
  96       try:
  97         handle, filepath = tempfile.mkstemp(prefix='isolateserver')
  98         # Write 2.1gb.
  99         chunk = chr(0) + chr(57) + chr(128) + chr(255)
 100         chunk1mb = chunk * (1024 * 1024 / len(chunk))
 101         for _ in xrange(1280):
 102           os.write(handle, chunk1mb)
 103       finally:
 104         os.close(handle)
 105
 106       self._archive_given_files([filepath])
 107     finally:
 108       if filepath:
 109         os.remove(filepath)
 110
 111
 112 if __name__ == '__main__':
 113   if len(sys.argv) > 1 and sys.argv[1].startswith('http'):
 114     ISOLATE_SERVER = sys.argv.pop(1).rstrip('/') + '/'
 115   logging.basicConfig(
 116       level=logging.DEBUG if '-v' in sys.argv else logging.ERROR)
 117   unittest.main()