Add asset management scripts

infra/bots/zip_utils_test.py

+#!/usr/bin/env python
+#
+# Copyright 2016 Google Inc.
+#
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+
+"""Tests for zip_utils."""
+
+
+import filecmp
+import os
+import unittest
+import utils
+import uuid
+import zip_utils
+
+
+def writefile(fname, mode=0640):
+  """Write the file with the given mode and random contents."""
+  dirname = os.path.dirname(fname)
+  if not os.path.isdir(dirname):
+    os.makedirs(dirname)
+  with open(fname, 'w') as f:
+    f.write(str(uuid.uuid4()))
+  os.chmod(fname, mode)
+
+
+class ZipUtilsTest(unittest.TestCase):
+  def _compare_trees(self, a, b):
+    """Compare two directory trees, assert if any differences."""
+    def _cmp(prefix, dcmp):
+      # Verify that the file and directory names are the same.
+      self.assertEqual(len(dcmp.left_only), 0)
+      self.assertEqual(len(dcmp.right_only), 0)
+      self.assertEqual(len(dcmp.diff_files), 0)
+      self.assertEqual(len(dcmp.funny_files), 0)
+
+      # Verify that the files are identical.
+      for f in dcmp.common_files:
+        pathA = os.path.join(a, prefix, f)
+        pathB = os.path.join(b, prefix, f)
+        self.assertTrue(filecmp.cmp(pathA, pathB, shallow=False))
+        statA = os.stat(pathA)
+        statB = os.stat(pathB)
+        self.assertEqual(statA.st_mode, statB.st_mode)
+        with open(pathA, 'rb') as f:
+          contentsA = f.read()
+        with open(pathB, 'rb') as f:
+          contentsB = f.read()
+        self.assertEqual(contentsA, contentsB)
+
+      # Recurse on subdirectories.
+      for prefix, obj in dcmp.subdirs.iteritems():
+        _cmp(prefix, obj)
+
+    _cmp('', filecmp.dircmp(a, b))
+
+  def test_zip_unzip(self):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      os.mkdir('input')
+      os.mkdir(os.path.join('input', 'mydir'))
+      os.mkdir(os.path.join('input', 'anotherdir'))
+      os.chmod(os.path.join('input', 'anotherdir'), 0666)
+      os.mkdir(os.path.join('input', 'dir3'))
+      os.chmod(os.path.join('input', 'dir3'), 0600)
+      writefile(os.path.join('input', 'a.txt'), 0777)
+      writefile(os.path.join('input', 'b.txt'), 0751)
+      writefile(os.path.join('input', 'c.txt'), 0640)
+      writefile(os.path.join('input', 'subdir', 'd.txt'), 0640)
+      
+      # Zip, unzip.
+      zip_utils.zip('input', 'test.zip')
+      zip_utils.unzip('test.zip', 'output')
+
+      # Compare the inputs and outputs.
+      self._compare_trees('input', 'output')
+
+  def test_blacklist(self):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      os.mkdir('input')
+      os.mkdir(os.path.join('input', '.git'))
+      writefile(os.path.join('input', '.git', 'index'))
+      writefile(os.path.join('input', 'somefile'))
+      writefile(os.path.join('input', '.DS_STORE'))
+      writefile(os.path.join('input', 'leftover.pyc'))
+      writefile(os.path.join('input', '.pycfile'))
+
+      # Zip, unzip.
+      zip_utils.zip('input', 'test.zip', blacklist=['.git', '.DS*', '*.pyc'])
+      zip_utils.unzip('test.zip', 'output')
+
+      # Remove the files/dirs we don't expect to see in output, so that we can
+      # use self._compare_trees to check the results.
+      os.remove(os.path.join('input', '.git', 'index'))
+      os.rmdir(os.path.join('input', '.git'))
+      os.remove(os.path.join('input', '.DS_STORE'))
+      os.remove(os.path.join('input', 'leftover.pyc'))
+
+      # Compare results.
+      self._compare_trees('input', 'output')
+
+
+if __name__ == '__main__':
+  unittest.main()