Client side implementation of new /content-gs isolate protocol.

tests/isolateserver_test.py

 #!/usr/bin/env python
 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 # pylint: disable=W0223
 # pylint: disable=W0231
 
-import binascii
 import hashlib
 import json
 import logging
 import os
-import random
 import shutil
 import StringIO
 import sys
 import tempfile
 import threading
 import unittest
+import urllib
 import zlib
 
 BASE_PATH = os.path.dirname(os.path.abspath(__file__))
 ROOT_DIR = os.path.dirname(BASE_PATH)
 sys.path.insert(0, ROOT_DIR)
 
 import auto_stub
 import isolateserver
 
 from utils import threading_utils
@@ skipping 26 matching lines @@
       for i, n in enumerate(self._requests):
         if n[0] == url:
           _, expected_kwargs, result = self._requests.pop(i)
           self.assertEqual(expected_kwargs, kwargs)
           if result is not None:
             return isolateserver.net.HttpResponse.get_fake_response(result, url)
           return None
     self.fail('Unknown request %s' % url)
 
 
+class TestZipCompression(TestCase):
+  """Test zip_compress and zip_decompress generators."""
+
+  def test_compress_and_decompress(self):
+    """Test data === decompress(compress(data))."""
+    original = [str(x) for x in xrange(0, 1000)]
+    processed = isolateserver.zip_decompress(
+        isolateserver.zip_compress(original))
+    self.assertEqual(''.join(original), ''.join(processed))
+
+  def test_zip_bomb(self):
+    """Verify zip_decompress always returns small chunks."""
+    original = '\x00' * 100000
+    bomb = ''.join(isolateserver.zip_compress(original))
+    decompressed = []
+    chunk_size = 1000
+    for chunk in isolateserver.zip_decompress([bomb], chunk_size):
+      self.assertLessEqual(len(chunk), chunk_size)
+      decompressed.append(chunk)
+    self.assertEqual(original, ''.join(decompressed))
+
+  def test_bad_zip_file(self):
+    """Verify decompressing broken file raises IOError."""
+    with self.assertRaises(IOError):
+      ''.join(isolateserver.zip_decompress(['Im not a zip file']))
+
+
 class StorageTest(TestCase):
   """Tests for Storage methods."""
 
   @staticmethod
   def mock_push(side_effect=None):
     """Returns StorageApi subclass with mocked 'push' method."""
     class MockedStorageApi(isolateserver.StorageApi):
       def __init__(self):
         self.pushed = []
       def push(self, item, expected_size, content_generator, push_urls=None):
@@ skipping 206 matching lines @@
           if metadata['h'] == digest
       ]
       self.assertEqual(1, len(filenames))
       filename = filenames[0]
       data = files_data[filename]
       self.assertEqual(
           (digest, len(data), data, push_urls[filename]),
           push_call)
 
 
-class IsolateServerArchiveTest(TestCase):
-  def setUp(self):
-    super(IsolateServerArchiveTest, self).setUp()
-    self.mock(isolateserver, 'randomness', lambda: 'not_really_random')
-    self.mock(sys, 'stdout', StringIO.StringIO())
-
-  def test_present(self):
+class IsolateServerStorageApiTest(TestCase):
+  @staticmethod
+  def mock_handshake_request(server, token='fake token', error=None):
+    handshake_request = {
+      'protocol_version': isolateserver.ISOLATE_PROTOCOL_VERSION,
+      'client_app_version': isolateserver.__version__,
+      'fetcher': True,
+      'pusher': True,
+    }
+    handshake_response = {
+      'protocol_version': isolateserver.ISOLATE_PROTOCOL_VERSION,
+      'server_app_version': 'mocked server T1000',
+      'access_token': token,
+      'error': error,
+    }
+    return (
+      server + '/content-gs/handshake',
+      {
+        'content_type': 'application/json',
+        'method': 'POST',
+        'data': json.dumps(handshake_request, separators=(',', ':')),
+      },
+      json.dumps(handshake_response),
+    )
+
+  @staticmethod
+  def mock_fetch_request(server, namespace, item, data):
+    return (
+      server + '/content-gs/retrieve/%s/%s' % (namespace, item),
+      {'retry_404': True, 'read_timeout': 60},
+      data,
+    )
+
+  @staticmethod
+  def mock_contains_request(server, namespace, token, request, response):
+    url = server + '/content-gs/pre-upload/%s?token=%s' % (
+        namespace, urllib.quote(token))
+    return (
+      url,
+      {
+        'data': json.dumps(request, separators=(',', ':')),
+        'content_type': 'application/json',
+        'method': 'POST',
+      },
+      json.dumps(response),
+    )
+
+  def test_server_capabilities_success(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    access_token = 'fake token'
+    self._requests = [
+      self.mock_handshake_request(server, access_token),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    caps = storage.server_capabilities
+    self.assertEqual(access_token, caps['access_token'])
+
+  def test_server_capabilities_network_failure(self):
+    self.mock(isolateserver.net, 'url_open', lambda *_args, **_kwargs: None)
+    with self.assertRaises(isolateserver.MappingError):
+      storage = isolateserver.IsolateServer('http://example.com', 'default')
+      _ = storage.server_capabilities
+
+  def test_server_capabilities_format_failure(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    handshake_req = self.mock_handshake_request(server)
+    self._requests = [
+      (handshake_req[0], handshake_req[1], 'Im a bad response'),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(isolateserver.MappingError):
+      _ = storage.server_capabilities
+
+  def test_server_capabilities_respects_error(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    error = 'Im sorry, Dave. Im afraid I cant do that.'
+    self._requests = [
+      self.mock_handshake_request(server, error=error)
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(isolateserver.MappingError) as context:
+      _ = storage.server_capabilities
+    # Server error message should be reported to user.
+    self.assertIn(error, str(context.exception))
+
+  def test_fetch_success_default(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    data = ''.join(str(x) for x in xrange(1000))
+    item = ALGO(data).hexdigest()
+    self._requests = [
+      self.mock_fetch_request(server, namespace, item, data),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    fetched = ''.join(storage.fetch(item, len(data)))
+    self.assertEqual(data, fetched)
+
+  def test_fetch_success_default_gzip(self):
+    server = 'http://example.com'
+    namespace = 'default-gzip'
+    data = ''.join(str(x) for x in xrange(1000))
+    item = ALGO(data).hexdigest()
+    self._requests = [
+      self.mock_fetch_request(server, namespace, item, zlib.compress(data)),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    fetched = ''.join(storage.fetch(item, len(data)))
+    self.assertEqual(data, fetched)
+
+  def test_fetch_failure_missing(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    item = ALGO('something').hexdigest()
+    self._requests = [
+      self.mock_fetch_request(server, namespace, item, None),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(IOError):
+      _ = ''.join(storage.fetch(item, isolateserver.UNKNOWN_FILE_SIZE))
+
+  def test_fetch_failure_bad_size(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    data = ''.join(str(x) for x in xrange(1000))
+    expected_size = len(data)
+    item = ALGO(data).hexdigest()
+    self._requests = [
+      self.mock_fetch_request(server, namespace, item, data[:100]),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(IOError):
+      _ = ''.join(storage.fetch(item, expected_size))
+
+  def test_fetch_failure_bad_zip(self):
+    server = 'http://example.com'
+    namespace = 'default-gzip'
+    item = ALGO('something').hexdigest()
+    self._requests = [
+      self.mock_fetch_request(server, namespace, item, 'Im not a zip'),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(IOError):
+      _ = ''.join(storage.fetch(item, isolateserver.UNKNOWN_FILE_SIZE))
+
+  def test_push_success(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    data = ''.join(str(x) for x in xrange(1000))
+    item = ALGO(data).hexdigest()
+    push_urls = (server + '/push_here', server + '/call_this')
+    self._requests = [
+      (
+        push_urls[0],
+        {
+          'data': data,
+          'content_type': 'application/octet-stream',
+          'method': 'PUT',
+        },
+        ''
+      ),
+      (
+        push_urls[1],
+        {
+          'data': '',
+          'content_type': 'application/json',
+          'method': 'POST',
+        },
+        ''
+      ),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    storage.push(item, len(data), [data], push_urls)
+
+  def test_push_failure_upload(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    data = ''.join(str(x) for x in xrange(1000))
+    item = ALGO(data).hexdigest()
+    push_urls = (server + '/push_here', server + '/call_this')
+    self._requests = [
+      (
+        push_urls[0],
+        {
+          'data': data,
+          'content_type': 'application/octet-stream',
+          'method': 'PUT',
+        },
+        None
+      ),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(IOError):
+      storage.push(item, len(data), [data], push_urls)
+
+  def test_push_failure_finalize(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    data = ''.join(str(x) for x in xrange(1000))
+    item = ALGO(data).hexdigest()
+    push_urls = (server + '/push_here', server + '/call_this')
+    self._requests = [
+      (
+        push_urls[0],
+        {
+          'data': data,
+          'content_type': 'application/octet-stream',
+          'method': 'PUT',
+        },
+        ''
+      ),
+      (
+        push_urls[1],
+        {
+          'data': '',
+          'content_type': 'application/json',
+          'method': 'POST',
+        },
+        None
+      ),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(IOError):
+      storage.push(item, len(data), [data], push_urls)
+
+  def test_contains_success(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    token = 'fake token'
     files = [
-      os.path.join(BASE_PATH, 'isolateserver', f)
-      for f in ('small_file.txt', 'empty_file.txt')
-    ]
-    hash_encoded = ''.join(
-        binascii.unhexlify(isolateserver.hash_file(f, ALGO)) for f in files)
-    path = 'http://random/'
-    self._requests = [
-      (path + 'content/get_token', {}, 'foo bar'),
-      (
-        path + 'content/contains/default-gzip?token=foo%20bar',
-        {'data': hash_encoded, 'content_type': 'application/octet-stream'},
-        '\1\1',
-      ),
-    ]
-    result = isolateserver.main(['archive', '--isolate-server', path] + files)
-    self.assertEqual(0, result)
-
-  def test_missing(self):
-    files = [
-      os.path.join(BASE_PATH, 'isolateserver', f)
-      for f in ('small_file.txt', 'empty_file.txt')
-    ]
-    hashes = [isolateserver.hash_file(f, ALGO) for f in files]
-    hash_encoded = ''.join(map(binascii.unhexlify, hashes))
-    compressed = [
-        zlib.compress(
-            open(f, 'rb').read(),
-            isolateserver.get_zip_compression_level(f))
-        for f in files
-    ]
-    path = 'http://random/'
-    self._requests = [
-      (path + 'content/get_token', {}, 'foo bar'),
-      (
-        path + 'content/contains/default-gzip?token=foo%20bar',
-        {'data': hash_encoded, 'content_type': 'application/octet-stream'},
-        '\0\0',
-      ),
-      (
-        path + 'content/store/default-gzip/%s?token=foo%%20bar' % hashes[0],
-        {'data': compressed[0], 'content_type': 'application/octet-stream'},
-        'ok',
-      ),
-      (
-        path + 'content/store/default-gzip/%s?token=foo%%20bar' % hashes[1],
-        {'data': compressed[1], 'content_type': 'application/octet-stream'},
-        'ok',
-      ),
-    ]
-    result = isolateserver.main(['archive', '--isolate-server', path] + files)
-    self.assertEqual(0, result)
-
-  def test_large(self):
-    content = ''
-    compressed = ''
-    while (
-        len(compressed) <= isolateserver.MIN_SIZE_FOR_DIRECT_BLOBSTORE):
-      # The goal here is to generate a file, once compressed, is at least
-      # MIN_SIZE_FOR_DIRECT_BLOBSTORE.
-      content += ''.join(chr(random.randint(0, 255)) for _ in xrange(20*1024))
-      compressed = zlib.compress(
-          content, isolateserver.get_zip_compression_level('foo.txt'))
-
-    s = ALGO(content).hexdigest()
-    infiles = {
-      'foo.txt': {
-        's': len(content),
-        'h': s,
-      },
-    }
-    path = 'http://random/'
-    hash_encoded = binascii.unhexlify(s)
-    content_type, body = isolateserver.encode_multipart_formdata(
-                [('token', 'foo bar')], [('content', s, compressed)])
-
-    self._requests = [
-      (path + 'content/get_token', {}, 'foo bar'),
-      (
-        path + 'content/contains/default-gzip?token=foo%20bar',
-        {'data': hash_encoded, 'content_type': 'application/octet-stream'},
-        '\0',
-      ),
-      (
-        path + 'content/generate_blobstore_url/default-gzip/%s' % s,
-        {'data': [('token', 'foo bar')]},
-        'an_url/',
-      ),
-      (
-        'an_url/',
-        {'data': body, 'content_type': content_type, 'retry_50x': False},
-        'ok',
-      ),
-    ]
-
-    # Setup mocks for zip_compress to return |compressed|.
-    self.mock(isolateserver, 'file_read', lambda *_: None)
-    self.mock(isolateserver, 'zip_compress', lambda *_: [compressed])
-    result = isolateserver.upload_tree(
-          base_url=path,
-          indir=os.getcwd(),
-          infiles=infiles,
-          namespace='default-gzip')
-
-    self.assertEqual(0, result)
-
-  def test_upload_blobstore_simple(self):
-    # A tad over 20kb so it triggers uploading to the blob store.
-    content = '0123456789' * 21*1024
-    s = ALGO(content).hexdigest()
-    path = 'http://example.com:80/'
-    data = [('token', 'a_token')]
-    content_type, body = isolateserver.encode_multipart_formdata(
-        data, [('content', s, content)])
-    self._requests = [
-      (
-        path + 'content/get_token',
-        {},
-        'a_token',
-      ),
-      (
-        path + 'content/generate_blobstore_url/x/' + s,
-        {'data': data[:]},
-        'http://example.com/an_url/',
-      ),
-      (
-        'http://example.com/an_url/',
-        {'data': body, 'content_type': content_type, 'retry_50x': False},
-        'ok42',
-      ),
-    ]
-    # |size| is currently ignored.
-    result = isolateserver.IsolateServer(path, 'x').push(s, -2, [content])
-    self.assertEqual('ok42', result)
-
-  def test_upload_blobstore_retry_500(self):
-    # A tad over 20kb so it triggers uploading to the blob store.
-    content = '0123456789' * 21*1024
-    s = ALGO(content).hexdigest()
-    path = 'http://example.com:80/'
-    data = [('token', 'a_token')]
-    content_type, body = isolateserver.encode_multipart_formdata(
-        data, [('content', s, content)])
-    self._requests = [
-      (
-        path + 'content/get_token',
-        {},
-        'a_token',
-      ),
-      (
-        path + 'content/generate_blobstore_url/x/' + s,
-        {'data': data[:]},
-        'http://example.com/an_url/',
-      ),
-      (
-        'http://example.com/an_url/',
-        {'data': body, 'content_type': content_type, 'retry_50x': False},
-        # Let's say an HTTP 500 was returned.
-        None,
-      ),
-      # In that case, a new url must be generated since the last one may have
-      # been "consumed".
-      (
-        path + 'content/generate_blobstore_url/x/' + s,
-        {'data': data[:]},
-        'http://example.com/an_url_2/',
-      ),
-      (
-        'http://example.com/an_url_2/',
-        {'data': body, 'content_type': content_type, 'retry_50x': False},
-        'ok42',
-      ),
-    ]
-    # |size| is currently ignored.
-    result = isolateserver.IsolateServer(path, 'x').push(s, -2, [content])
-    self.assertEqual('ok42', result)
+      ('file_a', {'h': 'hash1', 's': 100}),
+      ('file_b', {'h': 'hash2', 's': 200, 'priority': '0'}),
+      ('file_c', {'h': 'hash3', 's': 300, 'priority': '255'}),
+    ]
+    request = [
+      {'h': 'hash1', 's': 100, 'i': 0},
+      {'h': 'hash2', 's': 200, 'i': 1},
+      {'h': 'hash3', 's': 300, 'i': 0},
+    ]
+    response = [
+      None,
+      ['http://example/upload_here_1', None],
+      ['http://example/upload_here_2', 'http://example/call_this'],
+    ]
+    missing = [
+      files[1] + (response[1],),
+      files[2] + (response[2],)
+    ]
+    self._requests = [
+      self.mock_handshake_request(server, token),
+      self.mock_contains_request(server, namespace, token, request, response),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    result = storage.contains(files)
+    self.assertEqual(missing, result)
+
+  def test_contains_network_failure(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    token = 'fake token'
+    req = self.mock_contains_request(server, namespace, token, [], [])
+    self._requests = [
+      self.mock_handshake_request(server, token),
+      (req[0], req[1], None),
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(isolateserver.MappingError):
+      storage.contains([])
+
+  def test_contains_format_failure(self):
+    server = 'http://example.com'
+    namespace = 'default'
+    token = 'fake token'
+    self._requests = [
+      self.mock_handshake_request(server, token),
+      self.mock_contains_request(server, namespace, token, [], [1, 2, 3])
+    ]
+    storage = isolateserver.IsolateServer(server, namespace)
+    with self.assertRaises(isolateserver.MappingError):
+      storage.contains([])
 
 
 class IsolateServerDownloadTest(TestCase):
   tempdir = None
 
   def tearDown(self):
     try:
       if self.tempdir:
         shutil.rmtree(self.tempdir)
     finally:
       super(IsolateServerDownloadTest, self).tearDown()
 
   def test_download_two_files(self):
     # Test downloading two files.
     actual = {}
     def out(key, generator):
       actual[key] = ''.join(generator)
     self.mock(isolateserver, 'file_write', out)
     server = 'http://example.com'
     self._requests = [
       (
-        server + '/content/retrieve/default-gzip/sha-1',
+        server + '/content-gs/retrieve/default-gzip/sha-1',
         {'read_timeout': 60, 'retry_404': True},
         zlib.compress('Coucou'),
       ),
       (
-        server + '/content/retrieve/default-gzip/sha-2',
+        server + '/content-gs/retrieve/default-gzip/sha-2',
         {'read_timeout': 60, 'retry_404': True},
         zlib.compress('Bye Bye'),
       ),
     ]
     cmd = [
       'download',
       '--isolate-server', server,
       '--target', ROOT_DIR,
       '--file', 'sha-1', 'path/to/a',
       '--file', 'sha-2', 'path/to/b',
@@ skipping 27 matching lines @@
       'files': dict(
           (k, {'h': ALGO(v).hexdigest(), 's': len(v)})
           for k, v in files.iteritems()),
     }
     isolated_data = json.dumps(isolated, sort_keys=True, separators=(',',':'))
     isolated_hash = ALGO(isolated_data).hexdigest()
     requests = [(v['h'], files[k]) for k, v in isolated['files'].iteritems()]
     requests.append((isolated_hash, isolated_data))
     self._requests = [
       (
-        server + '/content/retrieve/default-gzip/' + h,
+        server + '/content-gs/retrieve/default-gzip/' + h,
         {
           'read_timeout': isolateserver.DOWNLOAD_READ_TIMEOUT,
           'retry_404': True,
         },
         zlib.compress(v),
       ) for h, v in requests
     ]
     cmd = [
       'download',
       '--isolate-server', server,
@@ skipping 88 matching lines @@
     expected = gen_data(os.path.sep)
     self.assertEqual(expected, actual)
 
 
 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()