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 zlib
 
 BASE_PATH = os.path.dirname(os.path.abspath(__file__))
 ROOT_DIR = os.path.dirname(BASE_PATH)
@@ skipping 32 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 alwasy returns small chunks."""
+    chunk_size = 1000
+    bomb = ''.join(isolateserver.zip_compress('\x00' * 100000))
+    for chunk in isolateserver.zip_decompress([bomb], chunk_size):
+      self.assertLessEqual(len(chunk), chunk_size)
+
+
 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):
-    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)
-
-
 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()