luci-py/isolateserver.py: Add archive support when downloading.

client/tests/isolateserver_test.py

 #!/usr/bin/env python
 # Copyright 2013 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.
 
 # pylint: disable=W0212,W0223,W0231,W0613
 
 import base64
 import collections
 import hashlib
 import json
 import logging
+import io
 import os
 import StringIO
 import sys
 import tempfile
 import unittest
 import urllib
 import zlib
 
 # net_utils adjusts sys.path.
 import net_utils
@@ skipping 692 matching lines @@
     for item in items:
       pending.add(item.digest)
       queue.add(item.digest)
 
     # Wait for fetch to complete.
     while pending:
       fetched = queue.wait(pending)
       pending.discard(fetched)
 
     # Ensure fetched same data as was pushed.
-    self.assertEqual(
-        [i.buffer for i in items],
-        [cache.read(i.digest) for i in items])
+    actual = []
+    for i in items:
+      with cache.getfileobj(i.digest) as f:
+        actual.append(f.read())
+
+    self.assertEqual([i.buffer for i in items], actual)
 
   def test_push_and_fetch(self):
     self.run_push_and_fetch_test('default')
 
   def test_push_and_fetch_gzip(self):
     self.run_push_and_fetch_test('default-gzip')
 
   if sys.maxsize == (2**31) - 1:
     def test_archive_multiple_huge_file(self):
       self.server.discard_content()
@@ skipping 93 matching lines @@
       '--file', 'sha-1', 'path/to/a',
       '--file', 'sha-2', 'path/to/b',
     ]
     self.assertEqual(0, isolateserver.main(cmd))
     expected = {
       os.path.join(net_utils.ROOT_DIR, 'path/to/a'): 'Coucou',
       os.path.join(net_utils.ROOT_DIR, 'path/to/b'): 'Bye Bye',
     }
     self.assertEqual(expected, actual)
 
-  def test_download_isolated(self):
+  def test_download_isolated_simple(self):
     # Test downloading an isolated tree.
     actual = {}
-    def file_write_mock(key, generator):
-      actual[key] = ''.join(generator)
-    self.mock(isolateserver, 'file_write', file_write_mock)
+    def putfile_mock(srcfileobj, dstpath, file_mode=None):
+      actual[dstpath] = srcfileobj.read()
+    self.mock(isolateserver, 'putfile', putfile_mock)
     self.mock(os, 'makedirs', lambda _: None)
     server = 'http://example.com'
     files = {
       os.path.join('a', 'foo'): 'Content',
       'b': 'More content',
-      }
+    }
     isolated = {
       'command': ['Absurb', 'command'],
       'relative_cwd': 'a',
       'files': dict(
           (k, {'h': isolateserver_mock.hash_content(v), 's': len(v)})
           for k, v in files.iteritems()),
       'version': isolated_format.ISOLATED_FILE_VERSION,
     }
     isolated_data = json.dumps(isolated, sort_keys=True, separators=(',',':'))
     isolated_hash = isolateserver_mock.hash_content(isolated_data)
@@ skipping 26 matching lines @@
     self.expected_requests(requests)
     self.assertEqual(0, isolateserver.main(cmd))
     expected = dict(
         (os.path.join(self.tempdir, k), v) for k, v in files.iteritems())
     self.assertEqual(expected, actual)
     expected_stdout = (
         'To run this test please run from the directory %s:\n  Absurb command\n'
         % os.path.join(self.tempdir, 'a'))
     self.checkOutput(expected_stdout, '')
 
+  def test_download_isolated_archive(self):
+    # Test downloading an isolated tree.
+    actual = {}
+    def putfile_mock(srcfileobj, dstpath, file_mode=None, size=-1):
+      actual[dstpath] = srcfileobj.read(size)
+    self.mock(isolateserver, 'putfile', putfile_mock)
+    self.mock(os, 'makedirs', lambda _: None)
+    server = 'http://example.com'
+
+    files = {
+      os.path.join('a', 'foo'): 'Content',
+      'b': 'More content',
+      'c': 'Even more content!',
+    }
+
+    archive = (
+      # ar file header
+      '!<arch>\n'
+      # File 1 -------------------------
+      # (16 bytes) filename len
+      '#1/5            '
+      # file metadata
+      '1447140471  1000  1000  100640  '
+      # (10 bytes) Data size
+      '12        '
+      # (2 bytes) File magic
+      '\x60\n'
+      # (5 bytes) File name
+      'a/foo'
+      # (7 bytes) File data
+      'Content'
+      # File 2 -------------------------
+      # (16 bytes) filename
+      'b               '
+      # file metadata
+      '1447140471  1000  1000  100640  '
+      # (12 bytes) Data size
+      '12        '
+      # (2 bytes) File magic
+      '\x60\n'
+      # (12 bytes) File data
+      'More content'
+      '')
+
+    isolated = {
+      'command': ['Absurb', 'command'],
+      'relative_cwd': 'a',
+      'files': {
+        'archive1': {
+          'h': isolateserver_mock.hash_content(archive),
+          's': len(archive),
+          't': 'smallfiles-archive',
+        },
+        'c': {
+          'h': isolateserver_mock.hash_content(files['c']),
+          's': len(files['c']),
+        },
+      },
+      'version': isolated_format.ISOLATED_FILE_VERSION,
+    }
+    isolated_data = json.dumps(isolated, sort_keys=True, separators=(',',':'))
+    isolated_hash = isolateserver_mock.hash_content(isolated_data)
+    requests = [
+      (isolated['files']['archive1']['h'], archive),
+      (isolated['files']['c']['h'], files['c']),
+    ]
+    requests.append((isolated_hash, isolated_data))
+    requests = [
+      (
+        server + '/_ah/api/isolateservice/v1/retrieve',
+        {
+            'data': {
+                'digest': h.encode('utf-8'),
+                'namespace': {
+                    'namespace': 'default-gzip',
+                    'digest_hash': 'sha-1',
+                    'compression': 'flate',
+                },
+                'offset': 0,
+            },
+            'read_timeout': 60,
+        },
+        {'content': base64.b64encode(zlib.compress(v))},
+      ) for h, v in requests
+    ]
+    cmd = [
+      'download',
+      '--isolate-server', server,
+      '--target', self.tempdir,
+      '--isolated', isolated_hash,
+    ]
+    self.expected_requests(requests)
+    self.assertEqual(0, isolateserver.main(cmd))
+    expected = dict(
+        (os.path.join(self.tempdir, k), v) for k, v in files.iteritems())
+    self.assertEqual(expected, actual)
+    expected_stdout = (
+        'To run this test please run from the directory %s:\n  Absurb command\n'
+        % os.path.join(self.tempdir, 'a'))
+    self.checkOutput(expected_stdout, '')
+
 
+
 def get_storage(_isolate_server, namespace):
   class StorageFake(object):
     def __enter__(self, *_):
       return self
 
     def __exit__(self, *_):
       pass
 
     @property
     def hash_algo(self):  # pylint: disable=R0201
@@ skipping 171 matching lines @@
 
 
 if __name__ == '__main__':
   fix_encoding.fix_encoding()
   if '-v' in sys.argv:
     unittest.TestCase.maxDiff = None
   logging.basicConfig(
       level=(logging.DEBUG if '-v' in sys.argv else logging.CRITICAL))
   clear_env_vars()
   unittest.main()