Update cbuildbot.py to upload prebuilts from preflight buildbot.

prebuilt_unittest.py

Index: prebuilt_unittest.py
diff --git a/prebuilt_unittest.py b/prebuilt_unittest.py
index fe2387fbf71da63bf58d372869a8d4294f7ed313..5d4e106ba66e2835a405d6091e1566896519911e 100755
--- a/prebuilt_unittest.py
+++ b/prebuilt_unittest.py
@@ -3,6 +3,7 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import copy
 import mox
 import os
 import prebuilt
@@ -10,6 +11,21 @@ import shutil
 import tempfile
 import unittest
 from chromite.lib import cros_build_lib
+from chromite.lib.binpkg import PackageIndex
+
+PUBLIC_PACKAGES = [{'CPV': 'public1', 'SHA1': '1'},
+                   {'CPV': 'public2', 'SHA1': '2', 'PATH': 'foo.tgz'}]
+PRIVATE_PACKAGES = [{'CPV': 'private', 'SHA1': '3'}]
+
+
+def SimplePackageIndex(header=True, packages=True):
+   pkgindex = PackageIndex()
+   if header:
+     pkgindex.header['URI'] = 'http://www.example.com'
+   if packages:
+     pkgindex.packages = copy.deepcopy(PUBLIC_PACKAGES + PRIVATE_PACKAGES)
+   return pkgindex
+
 
 class TestUpdateFile(unittest.TestCase):
 
@@ -137,14 +153,6 @@ class TestPrebuiltFilters(unittest.TestCase):
 
 
 class TestPrebuilt(unittest.TestCase):
-  fake_path = '/b/cbuild/build/chroot/build/x86-dogfood/'
-  bin_package_mock = ['packages/x11-misc/shared-mime-info-0.70.tbz2',
-                      'packages/x11-misc/util-macros-1.5.0.tbz2',
-                      'packages/x11-misc/xbitmaps-1.1.0.tbz2',
-                      'packages/x11-misc/read-edid-1.4.2.tbz2',
-                      'packages/x11-misc/xdg-utils-1.0.2-r3.tbz2']
-
-  files_to_sync = [os.path.join(fake_path, file) for file in bin_package_mock]
 
   def setUp(self):
     self.mox = mox.Mox()
@@ -153,23 +161,17 @@ class TestPrebuilt(unittest.TestCase):
     self.mox.UnsetStubs()
     self.mox.VerifyAll()
 
-  def _generate_dict_results(self, gs_bucket_path):
-    """
-    Generate a dictionary result similar to GenerateUploadDict
-    """
-    results = {}
-    for entry in self.files_to_sync:
-      results[entry] = os.path.join(
-        gs_bucket_path, entry.replace(self.fake_path, '').lstrip('/'))
-    return results
-
   def testGenerateUploadDict(self):
+    base_local_path = '/b/cbuild/build/chroot/build/x86-dogfood/'
     gs_bucket_path = 'gs://chromeos-prebuilt/host/version'
-    self.mox.StubOutWithMock(cros_build_lib, 'ListFiles')
-    cros_build_lib.ListFiles(self.fake_path).AndReturn(self.files_to_sync)
+    local_path = os.path.join(base_local_path, 'public1.tbz2')
+    self.mox.StubOutWithMock(prebuilt.os.path, 'exists')
+    prebuilt.os.path.exists(local_path).AndReturn(True)
     self.mox.ReplayAll()
-    result = prebuilt.GenerateUploadDict(self.fake_path, gs_bucket_path)
-    self.assertEqual(result, self._generate_dict_results(gs_bucket_path))
+    pkgs = [{ 'CPV': 'public1' }]
+    result = prebuilt.GenerateUploadDict(base_local_path, gs_bucket_path, pkgs)
+    expected = { local_path: gs_bucket_path + '/public1.tbz2' }
+    self.assertEqual(result, expected)
 
   def testFailonUploadFail(self):
     """Make sure we fail if one of the upload processes fail."""
@@ -202,31 +204,113 @@ class TestPackagesFileFiltering(unittest.TestCase):
     self.mox.UnsetStubs()
     self.mox.VerifyAll()
 
-  def testFilterAllPackages(self):
+  def testFilterPkgIndex(self):
     self.mox.StubOutWithMock(prebuilt, 'ShouldFilterPackage')
-    prebuilt.ShouldFilterPackage("public1").AndReturn(False)
-    prebuilt.ShouldFilterPackage("private").AndReturn(True)
-    prebuilt.ShouldFilterPackage("public2").AndReturn(False)
-    full_packages_file = [
-      "foo: bar\n", "\n",
-      "CPV: public1\n", "foo: bar1\n", "\n",
-      "CPV: private\n", "foo: bar2\n", "\n",
-      "CPV: public2\n", "foo: bar3\n", "\n",
-    ]
-    private_packages_file = [
-      "foo: bar\n", "\n",
-      "CPV: public1\n", "foo: bar1\n", "\n",
-      "CPV: public2\n", "foo: bar3\n", "\n",
-    ]
+    prebuilt.ShouldFilterPackage('public1').AndReturn(False)
+    prebuilt.ShouldFilterPackage('public2').AndReturn(False)
+    prebuilt.ShouldFilterPackage('private').AndReturn(True)
+    pkgindex = SimplePackageIndex()
+    self.mox.ReplayAll()
+    prebuilt.FilterPackagesIndex(pkgindex)
+    self.assertEqual(pkgindex.packages, PUBLIC_PACKAGES)
+    self.assertEqual(pkgindex.modified, True)
+
+
+class TestGrabPackageURLs(unittest.TestCase):
+
+  def testEmptyIndex(self):
+    pkgindex = SimplePackageIndex(packages=False)
+    db = prebuilt.GrabPackageURLs(pkgindex)
+    self.assertEqual(db, {})
+
+  def testNormalIndex(self):
+    pkgindex = SimplePackageIndex()
+    db = prebuilt.GrabPackageURLs(pkgindex)
+    self.assertEqual(len(db), 3)
+    self.assertEqual(db['1'], 'http://www.example.com/public1.tbz2')
+    self.assertEqual(db['2'], 'http://www.example.com/foo.tgz')
+    self.assertEqual(db['3'], 'http://www.example.com/private.tbz2')
+
+  def testFailedGrabPackageURLs(self):
+    pkgindex = SimplePackageIndex(header=False)
+    self.assertRaises(KeyError, prebuilt.GrabPackageURLs, pkgindex)
+    pkgindex = SimplePackageIndex()
+    del pkgindex.packages[0]['CPV']
+    self.assertRaises(KeyError, prebuilt.GrabPackageURLs, pkgindex)
+    pkgindex = SimplePackageIndex()
+    del pkgindex.packages[0]['SHA1']
+    self.assertRaises(KeyError, prebuilt.GrabPackageURLs, pkgindex)
+
+
+class TestFindDuplicates(unittest.TestCase):
+
+  def testEmptyIndex(self):
+    pkgindex = SimplePackageIndex(packages=False)
+    dupdb = {}
+    orig, dups = prebuilt.FindDuplicates(pkgindex, dupdb)
+    self.assertEqual(orig, [])
+    self.assertEqual(dups, [])
+
+  def testNormalIndex(self):
+    pkgindex = SimplePackageIndex()
+    dupdb = {}
+    orig, dups = prebuilt.FindDuplicates(pkgindex, dupdb)
+    self.assertEqual(orig, pkgindex.packages)
+    self.assertEqual(dups, [])
+
+  def testDuplicateIndex(self):
+    pkgindex = SimplePackageIndex()
+    dupdb = prebuilt.GrabPackageURLs(pkgindex)
+    orig, dups = prebuilt.FindDuplicates(pkgindex, dupdb)
+    self.assertEqual(orig, [])
+    expected = [(pkg, dupdb[pkg['SHA1']]) for pkg in pkgindex.packages]
+    self.assertEqual(dups, expected)
+
+  def testPartialIndex(self):
+    pkgindex = SimplePackageIndex()
+    dupdb = prebuilt.GrabPackageURLs(pkgindex)
+    del dupdb['3']
+    orig, dups = prebuilt.FindDuplicates(pkgindex, dupdb)
+    self.assertEqual(orig, [pkgindex.packages[2]])
+    expected = [(pkg, dupdb[pkg['SHA1']]) for pkg in pkgindex.packages[:-1]]
+    self.assertEqual(dups, expected)
+
+
+class TestResolveDuplicates(unittest.TestCase):
+
+  def testNoDuplicates(self):
+    pkgindex = SimplePackageIndex()
+    prebuilt.ResolveDuplicates(pkgindex, [])
+    self.assertEqual(pkgindex.modified, False)
+    self.assertEqual(pkgindex.packages, SimplePackageIndex().packages)
+
+  def testDuplicates(self):
+    pkgindex = SimplePackageIndex()
+    dupdb = prebuilt.GrabPackageURLs(pkgindex)
+    _, dups = prebuilt.FindDuplicates(pkgindex, dupdb)
+    expected = copy.deepcopy(pkgindex.packages)
+    for pkg in expected:
+      pkg['PATH'] = os.path.basename(dupdb[pkg['SHA1']])
+    prebuilt.ResolveDuplicates(pkgindex, dups)
+    self.assertEqual(pkgindex.packages, expected)
+
+
+class TestWritePackageIndex(unittest.TestCase):
+
+  def setUp(self):
+    self.mox = mox.Mox()
+
+  def tearDown(self):
+    self.mox.UnsetStubs()
+    self.mox.VerifyAll()
+
+  def testSimple(self):
+    pkgindex = SimplePackageIndex()
+    self.mox.StubOutWithMock(pkgindex, 'Write')
+    pkgindex.Write(mox.IgnoreArg())
     self.mox.ReplayAll()
-    temp_packages_file = tempfile.NamedTemporaryFile()
-    temp_packages_file.write("".join(full_packages_file))
-    temp_packages_file.flush()
-    new_packages_file = prebuilt.FilterPackagesFile(temp_packages_file.name)
-    new_contents = open(new_packages_file.name).read()
-    self.assertEqual("".join(private_packages_file), new_contents)
-    self.assertEqual("".join(private_packages_file), new_packages_file.read())
-    new_packages_file.close()
+    f = prebuilt.WritePackageIndex(pkgindex)
+    self.assertEqual(f.read(), '')
 
 
 if __name__ == '__main__':