luci-py: Refactor file writing code to allow file objects.

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 121 matching lines @@
 
   def contains(self, items):
     self.contains_calls.append(items)
     missing = {}
     for item in items:
       if item.digest in self.missing_hashes:
         missing[item] = self.missing_hashes[item.digest]
     return missing
 
 
+class UtilsTest(TestCase):
+  """Tests for helper methods in isolateserver file."""
+
+  def assertFile(self, path, contents):
+    self.assertTrue(fs.exists(path), 'File %s doesn\'t exist!' % path)
+    self.assertMultiLineEqual(contents, fs.open(path, 'rb').read())
+
+  def test_file_read(self):
+    # TODO(maruel): Write test for file_read generator (or remove it).
+    pass
+
+  def test_file_write(self):
+    # TODO(maruel): Write test for file_write generator (or remove it).
+    pass
+
+  def test_fileobj_path(self):
+    # No path on in-memory objects
+    self.assertIs(None, isolateserver.fileobj_path(io.BytesIO('hello')))
+
+    # Path on opened files
+    thisfile = os.path.abspath(__file__.decode(sys.getfilesystemencoding()))
+    f = fs.open(thisfile)
+    result = isolateserver.fileobj_path(f)
+    self.assertIsInstance(result, unicode)
+    self.assertSequenceEqual(result, thisfile)
+
+    # Path on temporary files
+    tf = tempfile.NamedTemporaryFile()
+    result = isolateserver.fileobj_path(tf)
+    self.assertIsInstance(result, unicode)
+    self.assertSequenceEqual(result, tf.name)
+
+    # No path on files which are no longer on the file system
+    tf = tempfile.NamedTemporaryFile(delete=False)
+    fs.unlink(tf.name.decode(sys.getfilesystemencoding()))
+    self.assertIs(None, isolateserver.fileobj_path(tf))
+
+  def test_fileobj_copy_simple(self):
+    inobj = io.BytesIO('hello')
+    outobj = io.BytesIO()
+
+    isolateserver.fileobj_copy(outobj, inobj)
+    self.assertEqual('hello', outobj.getvalue())
+
+  def test_fileobj_copy_partial(self):
+    inobj = io.BytesIO('adatab')
+    outobj = io.BytesIO()
+    inobj.read(1)
+
+    isolateserver.fileobj_copy(outobj, inobj, size=4)
+    self.assertEqual('data', outobj.getvalue())
+
+  def test_fileobj_copy_partial_file_no_size(self):
+    with self.assertRaises(IOError):
+      inobj = io.BytesIO('hello')
+      outobj = io.BytesIO()
+
+      inobj.read(1)
+      isolateserver.fileobj_copy(outobj, inobj)
+
+  def test_fileobj_copy_size_but_file_short(self):
+    with self.assertRaises(IOError):
+      inobj = io.BytesIO('hello')
+      outobj = io.BytesIO()
+
+      isolateserver.fileobj_copy(outobj, inobj, size=10)
+
+  def test_putfile(self):
+    tmpoutdir = None
+    tmpindir = None
+
+    try:
+      tmpindir = tempfile.mkdtemp(prefix='isolateserver_test')
+      infile = os.path.join(tmpindir, u'in')
+      with fs.open(infile, 'wb') as f:
+        f.write('data')
+
+      tmpoutdir = tempfile.mkdtemp(prefix='isolateserver_test')
+
+      # Copy as fileobj
+      fo = os.path.join(tmpoutdir, u'fo')
+      isolateserver.putfile(io.BytesIO('data'), fo)
+      self.assertEqual(True, fs.exists(fo))
+      self.assertEqual(False, fs.islink(fo))
+      self.assertFile(fo, 'data')
+
+      # Copy with partial fileobj
+      pfo = os.path.join(tmpoutdir, u'pfo')
+      fobj = io.BytesIO('adatab')
+      fobj.read(1)  # Read the 'a'
+      isolateserver.putfile(fobj, pfo, size=4)
+      self.assertEqual(True, fs.exists(pfo))
+      self.assertEqual(False, fs.islink(pfo))
+      self.assertEqual('b', fobj.read())
+      self.assertFile(pfo, 'data')
+
+      # Copy as not readonly
+      cp = os.path.join(tmpoutdir, u'cp')
+      with fs.open(infile, 'rb') as f:
+        isolateserver.putfile(f, cp, file_mode=0755)
+      self.assertEqual(True, fs.exists(cp))
+      self.assertEqual(False, fs.islink(cp))
+      self.assertFile(cp, 'data')
+
+      # Use hardlink
+      hl = os.path.join(tmpoutdir, u'hl')
+      with fs.open(infile, 'rb') as f:
+        isolateserver.putfile(f, hl, use_symlink=False)
+      self.assertEqual(True, fs.exists(hl))
+      self.assertEqual(False, fs.islink(hl))
+      self.assertFile(hl, 'data')
+
+      # Use symlink
+      sl = os.path.join(tmpoutdir, u'sl')
+      with fs.open(infile, 'rb') as f:
+        isolateserver.putfile(f, sl, use_symlink=True)
+      self.assertEqual(True, fs.exists(sl))
+      self.assertEqual(True, fs.islink(sl))
+      self.assertEqual('data', fs.open(sl, 'rb').read())
+      self.assertFile(sl, 'data')
+
+    finally:
+      if tmpindir:
+        file_path.rmtree(tmpindir)
+      if tmpoutdir:
+        file_path.rmtree(tmpoutdir)
+
+
 class StorageTest(TestCase):
   """Tests for Storage methods."""
 
   def assertEqualIgnoringOrder(self, a, b):
     """Asserts that containers |a| and |b| contain same items."""
     self.assertEqual(len(a), len(b))
     self.assertEqual(set(a), set(b))
 
   def get_push_state(self, storage, item):
     missing = list(storage.get_missing_items([item]))
@@ skipping 552 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, size=-1, use_symlink=False):
+      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 152 matching lines @@
     return isolateserver.DiskCache(self.tempdir, self._policies, self._algo)
 
   def to_hash(self, content):
     return self._algo(content).hexdigest(), content
 
   def test_read_evict(self):
     self._free_disk = 1100
     h_a = self.to_hash('a')[0]
     with self.get_cache() as cache:
       cache.write(h_a, 'a')
-      self.assertEqual('a', cache.read(h_a))
+      with cache.getfileobj(h_a) as f:
+        self.assertEqual('a', f.read())
 
     with self.get_cache() as cache:
       cache.evict(h_a)
       with self.assertRaises(isolateserver.CacheMiss):
-        cache.read(h_a)
-
-  def test_link(self):
-    self._free_disk = 1100
-    cache = self.get_cache()
-    h_a = self.to_hash('a')[0]
-    cache.write(h_a, 'a')
-    mapped = tempfile.mkdtemp(prefix='isolateserver_test')
-    try:
-      cache.link(h_a, os.path.join(mapped, u'hl'), False, False)
-      cache.link(h_a, os.path.join(mapped, u'sl'), False, True)
-      self.assertEqual(sorted(['hl', 'sl']), sorted(os.listdir(mapped)))
-      self.assertEqual(False, fs.islink(os.path.join(mapped, u'hl')))
-      self.assertEqual(True, fs.islink(os.path.join(mapped, u'sl')))
-    finally:
-      file_path.rmtree(mapped)
+        cache.getfileobj(h_a)
 
   def test_policies_free_disk(self):
     with self.assertRaises(isolateserver.Error):
       self.get_cache().write(*self.to_hash('a'))
 
   def test_policies_fit(self):
     self._free_disk = 1100
     self.get_cache().write(*self.to_hash('a'*100))
 
   def test_policies_too_much(self):
@@ skipping 81 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()