Update cros_mark_as_stable.py to also update the private overlay

BUG=chromium-os:7218
TEST=./cros_mark_all_as_stable --tracking_branch=cros/master

Change-Id: I0f7c4e2d2aa4fc6325721901a710bfc33fa39a91

Review URL: http://codereview.chromium.org/3516025

cros_mark_as_stable.py

@@ -68,11 +68,31 @@ def _Print(message):
   if gflags.FLAGS.verbose:
     Info(message)
 
+
+def _BuildEBuildDictionary(overlays, package_list, commit_id_list):
+  for index in range(len(package_list)):
+    package = package_list[index]
+    commit_id = ''
+    if commit_id_list:
+      commit_id = commit_id_list[index]
+    ebuild = _EBuild(package, commit_id)
+    if ebuild.ebuild_path:
+      for overlay in overlays:
+        if ebuild.ebuild_path.startswith(overlay):
+          overlays[overlay].append(ebuild)
+          break
+      else:
+        Die('No overlay found for %s' % ebuild.ebuild_path)
+    else:
+      Die('No ebuild found for %s' % package)
+
+
 def _CheckOnStabilizingBranch():
   """Returns true if the git branch is on the stabilizing branch."""
   current_branch = _SimpleRunCommand('git branch | grep \*').split()[1]
   return current_branch == _STABLE_BRANCH_NAME
 
+
 def _CheckSaneArguments(package_list, commit_id_list, command):
   """Checks to make sure the flags are sane.  Dies if arguments are not sane."""
   if not command in _COMMAND_DICTIONARY.keys():
@@ -217,7 +237,7 @@ class _EBuild(object):
     path = _SimpleRunCommand(equery_cmd)
     if path:
       _Print('Unstable ebuild found at %s' % path)
-    return path
+    return path.rstrip()
 
   @classmethod
   def _ParseEBuildPath(cls, ebuild_path):
@@ -276,8 +296,10 @@ class EBuildStableMarker(object):
                                          self._ebuild.current_revision + 1)
 
     _Print('Creating new stable ebuild %s' % new_ebuild_path)
-    shutil.copyfile('%s-9999.ebuild' % self._ebuild.ebuild_path_no_version,
+    workon_ebuild = '%s-9999.ebuild' % self._ebuild.ebuild_path_no_version
-                    new_ebuild_path)
+    if not os.path.exists(workon_ebuild):
+      Die('Missing 9999 ebuild: %s' % workon_ebuild)
+    shutil.copyfile(workon_ebuild, new_ebuild_path)
 
     for line in fileinput.input(new_ebuild_path, inplace=1):
       # Has to be done here to get changes to sys.stdout from fileinput.input.
@@ -336,41 +358,38 @@ def main(argv):
     commit_id_list = None
   _CheckSaneArguments(package_list, commit_id_list, command)
 
-  overlay_directory = '%s/third_party/chromiumos-overlay' % gflags.FLAGS.srcroot
+  overlays = {
+    '%s/private-overlays/chromeos-overlay' % gflags.FLAGS.srcroot: [],
+    '%s/third_party/chromiumos-overlay' % gflags.FLAGS.srcroot: []
+  }
+  _BuildEBuildDictionary(overlays, package_list, commit_id_list)
 
-  os.chdir(overlay_directory)
+  for overlay, ebuilds in overlays.items():
+    if not os.path.exists(overlay):
+      continue
+    os.chdir(overlay)
 
     if command == 'clean':
       _Clean()
-  elif command == 'commit':
+    elif command == 'push':
+      _PushChange()
+    elif command == 'commit' and ebuilds:
       work_branch = _GitBranch(_STABLE_BRANCH_NAME)
       work_branch.CreateBranch()
       if not work_branch.Exists():
-      Die('Unable to create stabilizing branch in %s' %
+        Die('Unable to create stabilizing branch in %s' % overlay)
-                               overlay_directory)
+      for ebuild in ebuilds:
-    index = 0
         try:
-      for index in range(len(package_list)):
+          _Print('Working on %s' % ebuild.package)
-        # Gather the package and optional commit id to work on.
+          worker = EBuildStableMarker(ebuild)
-        package = package_list[index]
+          worker.RevEBuild(ebuild.commit_id)
-        commit_id = ""
+          message = _GIT_COMMIT_MESSAGE % (ebuild.package, ebuild.commit_id)
-        if commit_id_list:
+          worker.CommitChange(message)
-          commit_id = commit_id_list[index]
+        except (OSError, IOError):
-
+          Warning('Cannot rev %s\n' % ebuild.package,
-        _Print('Working on %s' % package)
+                  'Note you will have to go into %s '
-        worker = EBuildStableMarker(_EBuild(package, commit_id))
+                  'and reset the git repo yourself.' % overlay)
-        worker.RevEBuild(commit_id)
+          raise
-        worker.CommitChange(_GIT_COMMIT_MESSAGE % (package, commit_id))
-
-    except (OSError, IOError), e:
-      Warning('An exception occurred\n'
-              'Only the following packages were revved: %s\n'
-              'Note you will have to go into %s'
-              'and reset the git repo yourself.' %
-              (package_list[:index], overlay_directory))
-      raise e
-  elif command == 'push':
-    _PushChange()
 
 
 if __name__ == '__main__':

cros_mark_as_stable_unittest.py

@@ -157,6 +157,7 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
 
   def testRevEBuild(self):
     self.mox.StubOutWithMock(cros_mark_as_stable.fileinput, 'input')
+    self.mox.StubOutWithMock(cros_mark_as_stable.os.path, 'exists')
     self.mox.StubOutWithMock(cros_mark_as_stable.shutil, 'copyfile')
     m_file = self.mox.CreateMock(file)
 
@@ -165,9 +166,39 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
     mock_file = ['EAPI=2', 'CROS_WORKON_COMMIT=old_id',
                  'KEYWORDS=\"~x86 ~arm\"', 'src_unpack(){}']
 
-    cros_mark_as_stable.shutil.copyfile(
+    ebuild_9999 = self.m_ebuild.ebuild_path_no_version + '-9999.ebuild'
-        self.m_ebuild.ebuild_path_no_version + '-9999.ebuild',
+    cros_mark_as_stable.os.path.exists(ebuild_9999).AndReturn(True)
-        self.revved_ebuild_path)
+    cros_mark_as_stable.shutil.copyfile(ebuild_9999, self.revved_ebuild_path)
+    cros_mark_as_stable.fileinput.input(self.revved_ebuild_path,
+                                        inplace=1).AndReturn(mock_file)
+    m_file.write('EAPI=2')
+    m_file.write('CROS_WORKON_COMMIT="my_id"\n')
+    m_file.write('KEYWORDS="x86 arm"')
+    m_file.write('src_unpack(){}')
+    cros_mark_as_stable._SimpleRunCommand('git add ' + self.revved_ebuild_path)
+    cros_mark_as_stable._SimpleRunCommand('git rm ' + self.m_ebuild.ebuild_path)
+
+    self.mox.ReplayAll()
+    marker = cros_mark_as_stable.EBuildStableMarker(self.m_ebuild)
+    marker.RevEBuild('my_id', redirect_file=m_file)
+    self.mox.VerifyAll()
+
+  def testRevMissingEBuild(self):
+    self.mox.StubOutWithMock(cros_mark_as_stable.fileinput, 'input')
+    self.mox.StubOutWithMock(cros_mark_as_stable.os.path, 'exists')
+    self.mox.StubOutWithMock(cros_mark_as_stable.shutil, 'copyfile')
+    self.mox.StubOutWithMock(cros_mark_as_stable, 'Die')
+    m_file = self.mox.CreateMock(file)
+
+    # Prepare mock fileinput.  This tests to make sure both the commit id
+    # and keywords are changed correctly.
+    mock_file = ['EAPI=2', 'CROS_WORKON_COMMIT=old_id',
+                 'KEYWORDS=\"~x86 ~arm\"', 'src_unpack(){}']
+
+    ebuild_9999 = self.m_ebuild.ebuild_path_no_version + '-9999.ebuild'
+    cros_mark_as_stable.os.path.exists(ebuild_9999).AndReturn(False)
+    cros_mark_as_stable.Die("Missing 9999 ebuild: %s" % ebuild_9999)
+    cros_mark_as_stable.shutil.copyfile(ebuild_9999, self.revved_ebuild_path)
     cros_mark_as_stable.fileinput.input(self.revved_ebuild_path,
                                         inplace=1).AndReturn(mock_file)
     m_file.write('EAPI=2')
@@ -200,6 +231,51 @@ class EBuildStableMarkerTest(mox.MoxTestBase):
     #self.mox.VerifyAll()
     pass
 
+class BuildEBuildDictionaryTest(mox.MoxTestBase):
+
+  def setUp(self):
+    mox.MoxTestBase.setUp(self)
+    self.mox.StubOutWithMock(cros_mark_as_stable, '_SimpleRunCommand')
+    self.ebuild_path = '/path/test_package-0.0.1-r1.ebuild'
+    self.package = "test_package"
+
+  def testValidPackage(self):
+    overlays = {"/path": []}
+    cmd = ('ACCEPT_KEYWORDS="x86 arm amd64" '
+           'equery-x86-generic which %s 2> /dev/null' % self.package)
+    cros_mark_as_stable._SimpleRunCommand(cmd).AndReturn(self.ebuild_path)
+    self.mox.ReplayAll()
+    cros_mark_as_stable._BuildEBuildDictionary(overlays, [self.package], [])
+    self.assertEquals(len(overlays), 1)
+    self.assertEquals(overlays["/path"][0].package, self.package)
+    self.mox.VerifyAll()
+
+  def testPackageInDifferentOverlay(self):
+    self.mox.StubOutWithMock(cros_mark_as_stable, 'Die')
+    cros_mark_as_stable.Die("No overlay found for %s" % self.ebuild_path)
+    cmd = ('ACCEPT_KEYWORDS="x86 arm amd64" '
+           'equery-x86-generic which %s 2> /dev/null' % self.package)
+    cros_mark_as_stable._SimpleRunCommand(cmd).AndReturn(self.ebuild_path)
+    overlays = {"/newpath": []}
+    self.mox.ReplayAll()
+    cros_mark_as_stable._BuildEBuildDictionary(overlays, [self.package], [])
+    self.assertEquals(len(overlays), 1)
+    self.assertEquals(overlays["/newpath"], [])
+    self.mox.VerifyAll()
+
+  def testMissingPackage(self):
+    self.mox.StubOutWithMock(cros_mark_as_stable, 'Die')
+    cros_mark_as_stable.Die("No ebuild found for %s" % self.package)
+    cmd = ('ACCEPT_KEYWORDS="x86 arm amd64" '
+           'equery-x86-generic which %s 2> /dev/null' % self.package)
+    cros_mark_as_stable._SimpleRunCommand(cmd).AndReturn("")
+    self.mox.ReplayAll()
+    overlays = {"/path": []}
+    cros_mark_as_stable._BuildEBuildDictionary(overlays, [self.package], [])
+    self.assertEquals(len(overlays), 1)
+    self.assertEquals(overlays["/path"], [])
+    self.mox.VerifyAll()
+
 
 if __name__ == '__main__':
   unittest.main()