llvm_tools: llvm version detection for roll tools

This CL implements version detection for update_chromeos_llvm_hash so
that uprevs will automatically get the correct major version in the llvm
ebuild.

BUG=b:186802799
TEST=Updated unit tests and ran tools manually

Change-Id: Id3a21ff7f708ce38303e27202428fb3bb519bafe
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2910173
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Commit-Queue: Ryan Beltran <ryanbeltran@chromium.org>
Tested-by: Ryan Beltran <ryanbeltran@chromium.org>

4 files changed, 262 insertions, 56 deletions

diff --git a/llvm_tools/get_llvm_hash.py b/llvm_tools/get_llvm_hash.py
index 7c770e64..e6d36edf 100755
--- a/llvm_tools/get_llvm_hash.py
+++ b/llvm_tools/get_llvm_hash.py
@@ -9,7 +9,9 @@
 from __future__ import print_function
 
 import argparse
+import functools
 import os
+import re
 import shutil
 import subprocess
 import sys
@@ -63,6 +65,65 @@ def GetGitHashFrom(src_dir, version):
       git_llvm_rev.Rev(branch=git_llvm_rev.MAIN_BRANCH, number=version))
 
 
+def CheckoutBranch(src_dir, branch):
+  """Checks out and pulls from a branch in a git repo.
+
+  Args:
+    src_dir: The LLVM source tree.
+    branch: The git branch to checkout in src_dir.
+
+  Raises:
+    ValueError: Failed to checkout or pull branch version
+  """
+  CheckCommand(['git', '-C', src_dir, 'checkout', branch])
+  CheckCommand(['git', '-C', src_dir, 'pull'])
+
+
+def ParseLLVMMajorVersion(cmakelist):
+  """Reads CMakeList.txt file contents for LLVMMajor Version.
+
+  Args:
+    cmakelist: contents of CMakeList.txt
+
+  Returns:
+    The major version number as a string
+
+  Raises:
+    ValueError: The major version cannot be parsed from cmakelist
+  """
+  match = re.search(r'\n\s+set\(LLVM_VERSION_MAJOR (?P<major>\d+)\)', cmakelist)
+  if not match:
+    raise ValueError('Failed to parse CMakeList for llvm major version')
+  return match.group('major')
+
+
+@functools.lru_cache(maxsize=1)
+def GetLLVMMajorVersion(git_hash=None):
+  """Reads llvm/CMakeList.txt file contents for LLVMMajor Version.
+
+  Args:
+    git_hash: git hash of llvm version as string or None for top of trunk
+
+  Returns:
+    The major version number as a string
+
+  Raises:
+    ValueError: The major version cannot be parsed from cmakelist or
+      there was a failure to checkout git_hash version
+    FileExistsError: The src directory doe not contain CMakeList.txt
+  """
+  src_dir = GetAndUpdateLLVMProjectInLLVMTools()
+  cmakelists_path = os.path.join(src_dir, 'llvm', 'CMakeLists.txt')
+  if git_hash:
+    CheckCommand(['git', '-C', src_dir, 'checkout', git_hash])
+  try:
+    with open(cmakelists_path) as cmakelists_file:
+      return ParseLLVMMajorVersion(cmakelists_file.read())
+  finally:
+    if git_hash:
+      CheckoutBranch(src_dir, git_llvm_rev.MAIN_BRANCH)
+
+
 @contextmanager
 def CreateTempLLVMRepo(temp_dir):
   """Adds a LLVM worktree to 'temp_dir'.
@@ -125,9 +186,10 @@ def GetAndUpdateLLVMProjectInLLVMTools():
                                               'llvm-project-copy')
 
   if not os.path.isdir(abs_path_to_llvm_project_dir):
-    print('Checking out LLVM from scratch. This could take a while...\n'
-          '(This should only need to be done once, though.)',
-          file=sys.stderr)
+    print(
+        'Checking out LLVM from scratch. This could take a while...\n'
+        '(This should only need to be done once, though.)',
+        file=sys.stderr)
     os.mkdir(abs_path_to_llvm_project_dir)
 
     LLVMHash().CloneLLVMRepo(abs_path_to_llvm_project_dir)
@@ -142,11 +204,7 @@ def GetAndUpdateLLVMProjectInLLVMTools():
       raise ValueError('LLVM repo in %s has changes, please remove.' %
                        abs_path_to_llvm_project_dir)
 
-    CheckCommand([
-        'git', '-C', abs_path_to_llvm_project_dir, 'checkout',
-        git_llvm_rev.MAIN_BRANCH
-    ])
-    CheckCommand(['git', '-C', abs_path_to_llvm_project_dir, 'pull'])
+    CheckoutBranch(abs_path_to_llvm_project_dir, git_llvm_rev.MAIN_BRANCH)
 
   return abs_path_to_llvm_project_dir
 
diff --git a/llvm_tools/get_llvm_hash_unittest.py b/llvm_tools/get_llvm_hash_unittest.py
index 2e56aed5..49740f33 100755
--- a/llvm_tools/get_llvm_hash_unittest.py
+++ b/llvm_tools/get_llvm_hash_unittest.py
@@ -90,6 +90,50 @@ class TestGetLLVMHash(unittest.TestCase):
     self.assertEqual(LLVMHash().GetTopOfTrunkGitHash(), 'a123testhash1')
     mock_check_output.assert_called_once()
 
+  @mock.patch.object(subprocess, 'Popen')
+  def testCheckoutBranch(self, mock_popen):
+    mock_popen.return_value = mock.MagicMock(
+        communicate=lambda: (None, None), returncode=0)
+    get_llvm_hash.CheckoutBranch('fake/src_dir', 'fake_branch')
+    self.assertEqual(
+        mock_popen.call_args_list[0][0],
+        (['git', '-C', 'fake/src_dir', 'checkout', 'fake_branch'],))
+    self.assertEqual(mock_popen.call_args_list[1][0],
+                     (['git', '-C', 'fake/src_dir', 'pull'],))
+
+  def testParseLLVMMajorVersion(self):
+    cmakelist_42 = ('set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)\n'
+                    'if(NOT DEFINED LLVM_VERSION_MAJOR)\n'
+                    '  set(LLVM_VERSION_MAJOR 42)\n'
+                    'endif()')
+    self.assertEqual(get_llvm_hash.ParseLLVMMajorVersion(cmakelist_42), '42')
+
+  def testParseLLVMMajorVersionInvalid(self):
+    invalid_cmakelist = 'invalid cmakelist.txt contents'
+    with self.assertRaises(ValueError):
+      get_llvm_hash.ParseLLVMMajorVersion(invalid_cmakelist)
+
+  @mock.patch.object(get_llvm_hash, 'GetAndUpdateLLVMProjectInLLVMTools')
+  @mock.patch.object(get_llvm_hash, 'ParseLLVMMajorVersion')
+  @mock.patch.object(get_llvm_hash, 'CheckCommand')
+  @mock.patch.object(get_llvm_hash, 'CheckoutBranch')
+  @mock.patch(
+      'get_llvm_hash.open',
+      mock.mock_open(read_data='mock contents'),
+      create=True)
+  def testGetLLVMMajorVersion(self, mock_checkout_branch, mock_git_checkout,
+                              mock_major_version, mock_llvm_project_path):
+    mock_llvm_project_path.return_value = 'path/to/llvm-project'
+    mock_major_version.return_value = '1234'
+    self.assertEqual(get_llvm_hash.GetLLVMMajorVersion('314159265'), '1234')
+    # Second call should be memoized
+    self.assertEqual(get_llvm_hash.GetLLVMMajorVersion('314159265'), '1234')
+    mock_llvm_project_path.assert_called_once()
+    mock_major_version.assert_called_with('mock contents')
+    mock_git_checkout.assert_called_once_with(
+        ['git', '-C', 'path/to/llvm-project', 'checkout', '314159265'])
+    mock_checkout_branch.assert_called_once_with('path/to/llvm-project', 'main')
+
 
 if __name__ == '__main__':
   unittest.main()
diff --git a/llvm_tools/update_chromeos_llvm_hash.py b/llvm_tools/update_chromeos_llvm_hash.py
index e28fe690..a5227cc9 100755
--- a/llvm_tools/update_chromeos_llvm_hash.py
+++ b/llvm_tools/update_chromeos_llvm_hash.py
@@ -199,7 +199,6 @@ def UpdateEbuildLLVMHash(ebuild_path, llvm_variant, git_hash, svn_version):
       for cur_line in ReplaceLLVMHash(ebuild_file, llvm_variant, git_hash,
                                       svn_version):
         temp_file.write(cur_line)
-
   os.rename(temp_ebuild_file, ebuild_path)
 
   # Get the path to the parent directory.
@@ -266,7 +265,7 @@ def UprevEbuildSymlink(symlink):
        os.path.dirname(symlink), 'mv', symlink, new_symlink])
 
 
-def UprevEbuildToVersion(symlink, svn_version):
+def UprevEbuildToVersion(symlink, svn_version, git_hash):
   """Uprevs the ebuild's revision number.
 
   Increases the revision number by 1 and stages the change in
@@ -275,9 +274,11 @@ def UprevEbuildToVersion(symlink, svn_version):
   Args:
     symlink: The absolute path of an ebuild symlink.
     svn_version: The SVN-style revision number of git_hash.
+    git_hash: The new git hash.
 
   Raises:
     ValueError: Failed to uprev the ebuild or failed to stage the changes.
+    AssertionError: No llvm version provided for an LLVM uprev
   """
 
   if not os.path.islink(symlink):
@@ -290,10 +291,11 @@ def UprevEbuildToVersion(symlink, svn_version):
     raise ValueError('Tried to uprev an unknown package')
   # llvm
   if package == 'llvm':
+    llvm_major_version = get_llvm_hash.GetLLVMMajorVersion(git_hash)
     new_ebuild, is_changed = re.subn(
-        r'pre([0-9]+)_p([0-9]+)',
-        'pre%s_p%s' % (svn_version, \
-            datetime.today().strftime('%Y%m%d')),
+        r'(\d+)\.(\d+)_pre([0-9]+)_p([0-9]+)',
+        '%s.\\2_pre%s_p%s' %
+        (llvm_major_version, svn_version, datetime.today().strftime('%Y%m%d')),
         ebuild,
         count=1)
   # any other package
@@ -374,8 +376,8 @@ def StagePatchMetadataFileForCommit(patch_metadata_file_path):
   """
 
   if not os.path.isfile(patch_metadata_file_path):
-    raise ValueError(
-        'Invalid patch metadata file provided: %s' % patch_metadata_file_path)
+    raise ValueError('Invalid patch metadata file provided: %s' %
+                     patch_metadata_file_path)
 
   # Cmd to stage the patch metadata file for commit.
   subprocess.check_output([
@@ -432,8 +434,8 @@ def StagePackagesPatchResultsForCommit(package_info_dict, commit_messages):
   return commit_messages
 
 
-def UpdatePackages(packages, llvm_variant, git_hash, svn_version, chroot_path,
-                   patch_metadata_file, mode, git_hash_source,
+def UpdatePackages(packages, llvm_variant, git_hash, svn_version,
+                   chroot_path, patch_metadata_file, mode, git_hash_source,
                    extra_commit_msg):
   """Updates an LLVM hash and uprevs the ebuild of the packages.
 
@@ -478,11 +480,11 @@ def UpdatePackages(packages, llvm_variant, git_hash, svn_version, chroot_path,
     if llvm_variant == LLVMVariant.next:
       commit_message_header = 'llvm-next'
     if git_hash_source in get_llvm_hash.KNOWN_HASH_SOURCES:
-      commit_message_header += (
-          '/%s: upgrade to %s (r%d)' % (git_hash_source, git_hash, svn_version))
+      commit_message_header += ('/%s: upgrade to %s (r%d)' %
+                                (git_hash_source, git_hash, svn_version))
     else:
-      commit_message_header += (
-          ': upgrade to %s (r%d)' % (git_hash, svn_version))
+      commit_message_header += (': upgrade to %s (r%d)' %
+                                (git_hash, svn_version))
 
     commit_messages = [
         commit_message_header + '\n',
@@ -504,7 +506,7 @@ def UpdatePackages(packages, llvm_variant, git_hash, svn_version, chroot_path,
       UpdateEbuildLLVMHash(ebuild_path, llvm_variant, git_hash, svn_version)
 
       if llvm_variant == LLVMVariant.current:
-        UprevEbuildToVersion(symlink_path, svn_version)
+        UprevEbuildToVersion(symlink_path, svn_version, git_hash)
       else:
         UprevEbuildSymlink(symlink_path)
 
@@ -514,6 +516,8 @@ def UpdatePackages(packages, llvm_variant, git_hash, svn_version, chroot_path,
       packages.append('%s/%s' % (parent_dir_name, cur_dir_name))
       commit_messages.append('%s/%s' % (parent_dir_name, cur_dir_name))
 
+    EnsurePackageMaskContains(chroot_path, git_hash)
+
     # Handle the patches for each package.
     package_info_dict = llvm_patch_management.UpdatePackagesPatchMetadataFile(
         chroot_path, svn_version, patch_metadata_file, packages, mode)
@@ -533,6 +537,31 @@ def UpdatePackages(packages, llvm_variant, git_hash, svn_version, chroot_path,
   return change_list
 
 
+def EnsurePackageMaskContains(chroot_path, git_hash):
+  """Adds the major version of llvm to package.mask if it's not already present.
+
+  Args:
+    chroot_path: The absolute path to the chroot.
+    git_hash: The new git hash.
+
+  Raises:
+    FileExistsError: package.mask not found in ../../chromiumos-overlay
+  """
+
+  llvm_major_version = get_llvm_hash.GetLLVMMajorVersion(git_hash)
+
+  overlay_dir = os.path.join(chroot_path, 'src/third_party/chromiumos-overlay')
+  mask_path = os.path.join(overlay_dir,
+                           'profiles/targets/chromeos/package.mask')
+  with open(mask_path, 'r+') as mask_file:
+    mask_contents = mask_file.read()
+    expected_line = '=sys-devel/llvm-%s.0_pre*\n' % llvm_major_version
+    if expected_line not in mask_contents:
+      mask_file.write(expected_line)
+
+  subprocess.check_output(['git', '-C', overlay_dir, 'add', mask_path])
+
+
 def main():
   """Updates the LLVM next hash for each package.
 
diff --git a/llvm_tools/update_chromeos_llvm_hash_unittest.py b/llvm_tools/update_chromeos_llvm_hash_unittest.py
index 205feb0c..2e928be9 100755
--- a/llvm_tools/update_chromeos_llvm_hash_unittest.py
+++ b/llvm_tools/update_chromeos_llvm_hash_unittest.py
@@ -18,6 +18,7 @@ import unittest.mock as mock
 
 import chroot
 import failure_modes
+import get_llvm_hash
 import git
 import llvm_patch_management
 import test_helpers
@@ -67,8 +68,9 @@ class UpdateLLVMHashTest(unittest.TestCase):
       # Verify the exception is raised when the ebuild file does not have
       # 'LLVM_HASH'.
       with self.assertRaises(ValueError) as err:
-        update_chromeos_llvm_hash.UpdateEbuildLLVMHash(
-            ebuild_file, llvm_variant, git_hash, svn_version)
+        update_chromeos_llvm_hash.UpdateEbuildLLVMHash(ebuild_file,
+                                                       llvm_variant, git_hash,
+                                                       svn_version)
 
         self.assertEqual(
             str(err.exception), ('Failed to update %s.', 'LLVM_HASH'))
@@ -95,8 +97,9 @@ class UpdateLLVMHashTest(unittest.TestCase):
       # Verify the exception is raised when the ebuild file does not have
       # 'LLVM_NEXT_HASH'.
       with self.assertRaises(ValueError) as err:
-        update_chromeos_llvm_hash.UpdateEbuildLLVMHash(
-            ebuild_file, llvm_variant, git_hash, svn_version)
+        update_chromeos_llvm_hash.UpdateEbuildLLVMHash(ebuild_file,
+                                                       llvm_variant, git_hash,
+                                                       svn_version)
 
         self.assertEqual(
             str(err.exception), ('Failed to update %s.', 'LLVM_NEXT_HASH'))
@@ -179,19 +182,25 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_stage_commit_command.assert_called_once()
 
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
   @mock.patch.object(os.path, 'islink', return_value=False)
-  def testFailedToUprevEbuildToVersionForInvalidSymlink(self, mock_islink):
+  def testFailedToUprevEbuildToVersionForInvalidSymlink(self, mock_islink,
+                                                        mock_llvm_version):
     symlink_path = '/path/to/chroot/package/package.ebuild'
     svn_version = 1000
+    git_hash = 'badf00d'
+    mock_llvm_version.return_value = '1234'
 
     # Verify the exception is raised when a invalid symbolic link is passed in.
     with self.assertRaises(ValueError) as err:
-      update_chromeos_llvm_hash.UprevEbuildToVersion(symlink_path, svn_version)
+      update_chromeos_llvm_hash.UprevEbuildToVersion(symlink_path, svn_version,
+                                                     git_hash)
 
     self.assertEqual(
         str(err.exception), 'Invalid symlink provided: %s' % symlink_path)
 
     mock_islink.assert_called_once()
+    mock_llvm_version.assert_not_called()
 
   @mock.patch.object(os.path, 'islink', return_value=False)
   def testFailedToUprevEbuildSymlinkForInvalidSymlink(self, mock_islink):
@@ -206,22 +215,28 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_islink.assert_called_once()
 
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
   # Simulate 'os.path.islink' when a symbolic link is passed in.
   @mock.patch.object(os.path, 'islink', return_value=True)
   # Simulate 'os.path.realpath' when a symbolic link is passed in.
   @mock.patch.object(os.path, 'realpath', return_value=True)
-  def testFailedToUprevEbuildToVersion(self, mock_realpath, mock_islink):
+  def testFailedToUprevEbuildToVersion(self, mock_realpath, mock_islink,
+                                       mock_llvm_version):
     symlink_path = '/path/to/chroot/llvm/llvm_pre123_p.ebuild'
     mock_realpath.return_value = '/abs/path/to/llvm/llvm_pre123_p.ebuild'
+    git_hash = 'badf00d'
+    mock_llvm_version.return_value = '1234'
     svn_version = 1000
 
     # Verify the exception is raised when the symlink does not match the
     # expected pattern
     with self.assertRaises(ValueError) as err:
-      update_chromeos_llvm_hash.UprevEbuildToVersion(symlink_path, svn_version)
+      update_chromeos_llvm_hash.UprevEbuildToVersion(symlink_path, svn_version,
+                                                     git_hash)
 
     self.assertEqual(str(err.exception), 'Failed to uprev the ebuild.')
 
+    mock_llvm_version.assert_called_once_with(git_hash)
     mock_islink.assert_called_once_with(symlink_path)
 
   # Simulate 'os.path.islink' when a symbolic link is passed in.
@@ -238,17 +253,24 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_islink.assert_called_once_with(symlink_path)
 
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
   @mock.patch.object(os.path, 'islink', return_value=True)
   @mock.patch.object(os.path, 'realpath')
   @mock.patch.object(subprocess, 'check_output', return_value=None)
   def testSuccessfullyUprevEbuildToVersionLLVM(self, mock_command_output,
-                                               mock_realpath, mock_islink):
-    symlink = '/path/to/llvm/llvm_pre3_p2-r10.ebuild'
-    ebuild = '/abs/path/to/llvm/llvm_pre3_p2.ebuild'
+                                               mock_realpath, mock_islink,
+                                               mock_llvm_version):
+    symlink = '/path/to/llvm/llvm-12.0_pre3_p2-r10.ebuild'
+    ebuild = '/abs/path/to/llvm/llvm-12.0_pre3_p2.ebuild'
     mock_realpath.return_value = ebuild
+    git_hash = 'badf00d'
+    mock_llvm_version.return_value = '1234'
     svn_version = 1000
 
-    update_chromeos_llvm_hash.UprevEbuildToVersion(symlink, svn_version)
+    update_chromeos_llvm_hash.UprevEbuildToVersion(symlink, svn_version,
+                                                   git_hash)
+
+    mock_llvm_version.assert_called_once_with(git_hash)
 
     mock_islink.assert_called()
 
@@ -258,12 +280,8 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     # Verify commands
     symlink_dir = os.path.dirname(symlink)
-    new_ebuild, _is_changed = re.subn(
-        r'pre([0-9]+)_p([0-9]+)',
-        'pre%s_p%s' % (svn_version, \
-            datetime.today().strftime('%Y%m%d')),
-        ebuild,
-        count=1)
+    timestamp = datetime.today().strftime('%Y%m%d')
+    new_ebuild = '/abs/path/to/llvm/llvm-1234.0_pre1000_p%s.ebuild' % timestamp
     new_symlink = new_ebuild[:-len('.ebuild')] + '-r1.ebuild'
 
     expected_cmd = ['git', '-C', symlink_dir, 'mv', ebuild, new_ebuild]
@@ -291,8 +309,10 @@ class UpdateLLVMHashTest(unittest.TestCase):
     ebuild = '/abs/path/to/compiler-rt/compiler-rt_pre3_p2.ebuild'
     mock_realpath.return_value = ebuild
     svn_version = 1000
+    git_hash = '1234'
 
-    update_chromeos_llvm_hash.UprevEbuildToVersion(symlink, svn_version)
+    update_chromeos_llvm_hash.UprevEbuildToVersion(symlink, svn_version,
+                                                   git_hash)
 
     mock_islink.assert_called()
 
@@ -549,6 +569,7 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     self.assertEqual(mock_stage_patches_for_commit.call_count, 2)
 
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
   @mock.patch.object(update_chromeos_llvm_hash,
                      'CreatePathDictionaryFromPackages')
   @mock.patch.object(git, 'CreateBranch')
@@ -560,7 +581,7 @@ class UpdateLLVMHashTest(unittest.TestCase):
   def testExceptionRaisedWhenUpdatingPackages(
       self, mock_realpath, mock_delete_repo, mock_upload_changes,
       mock_uprev_symlink, mock_update_llvm_next, mock_create_repo,
-      mock_create_path_dict):
+      mock_create_path_dict, mock_llvm_major_version):
 
     abs_path_to_package = '/some/path/to/chroot/src/path/to/package.ebuild'
 
@@ -569,6 +590,7 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     path_to_package_dir = '/some/path/to/chroot/src/path/to'
 
+    mock_llvm_major_version.return_value = '1234'
     # Test function to simulate 'CreateBranch' when successfully created the
     # branch on a valid repo path.
     def SuccessfullyCreateBranchForChanges(_repo_path, branch):
@@ -628,9 +650,9 @@ class UpdateLLVMHashTest(unittest.TestCase):
     # the 'try' block by UprevEbuildSymlink function.
     with self.assertRaises(ValueError) as err:
       update_chromeos_llvm_hash.UpdatePackages(
-          packages_to_update, llvm_variant, git_hash, svn_version, chroot_path,
-          patch_metadata_file, failure_modes.FailureModes.FAIL, git_hash_source,
-          extra_commit_msg)
+          packages_to_update, llvm_variant, git_hash, svn_version,
+          chroot_path, patch_metadata_file, failure_modes.FailureModes.FAIL,
+          git_hash_source, extra_commit_msg)
 
     self.assertEqual(str(err.exception), 'Failed to uprev the ebuild.')
 
@@ -639,8 +661,9 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_create_repo.assert_called_once_with(path_to_package_dir, branch)
 
-    mock_update_llvm_next.assert_called_once_with(
-        abs_path_to_package, llvm_variant, git_hash, svn_version)
+    mock_update_llvm_next.assert_called_once_with(abs_path_to_package,
+                                                  llvm_variant, git_hash,
+                                                  svn_version)
 
     mock_uprev_symlink.assert_called_once_with(symlink_path_to_package)
 
@@ -648,6 +671,8 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_delete_repo.assert_called_once_with(path_to_package_dir, branch)
 
+  @mock.patch.object(update_chromeos_llvm_hash, 'EnsurePackageMaskContains')
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
   @mock.patch.object(update_chromeos_llvm_hash,
                      'CreatePathDictionaryFromPackages')
   @mock.patch.object(git, 'CreateBranch')
@@ -658,10 +683,12 @@ class UpdateLLVMHashTest(unittest.TestCase):
   @mock.patch.object(llvm_patch_management, 'UpdatePackagesPatchMetadataFile')
   @mock.patch.object(update_chromeos_llvm_hash,
                      'StagePatchMetadataFileForCommit')
-  def testSuccessfullyUpdatedPackages(
-      self, mock_stage_patch_file, mock_update_package_metadata_file,
-      mock_delete_repo, mock_upload_changes, mock_uprev_symlink,
-      mock_update_llvm_next, mock_create_repo, mock_create_path_dict):
+  def testSuccessfullyUpdatedPackages(self, mock_stage_patch_file,
+                                      mock_update_package_metadata_file,
+                                      mock_delete_repo, mock_upload_changes,
+                                      mock_uprev_symlink, mock_update_llvm_next,
+                                      mock_create_repo, mock_create_path_dict,
+                                      mock_llvm_version, mock_mask_contains):
 
     abs_path_to_package = '/some/path/to/chroot/src/path/to/package.ebuild'
 
@@ -746,6 +773,8 @@ class UpdateLLVMHashTest(unittest.TestCase):
     mock_uprev_symlink.side_effect = SuccessfullyUprevedEbuildSymlink
     mock_update_package_metadata_file.side_effect = RetrievedPatchResults
     mock_upload_changes.side_effect = SuccessfullyUploadedChanges
+    mock_llvm_version.return_value = '1234'
+    mock_mask_contains.reurn_value = None
 
     packages_to_update = ['test-packages/package1']
     llvm_variant = update_chromeos_llvm_hash.LLVMVariant.next
@@ -758,9 +787,10 @@ class UpdateLLVMHashTest(unittest.TestCase):
     extra_commit_msg = '\ncommit-message-end'
 
     change_list = update_chromeos_llvm_hash.UpdatePackages(
-        packages_to_update, llvm_variant, git_hash, svn_version, chroot_path,
-        patch_metadata_file, failure_modes.FailureModes.DISABLE_PATCHES,
-        git_hash_source, extra_commit_msg)
+        packages_to_update, llvm_variant, git_hash, svn_version,
+        chroot_path, patch_metadata_file,
+        failure_modes.FailureModes.DISABLE_PATCHES, git_hash_source,
+        extra_commit_msg)
 
     self.assertEqual(change_list.url,
                      'https://some_name/path/to/commit/+/12345')
@@ -772,11 +802,14 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_create_repo.assert_called_once_with(path_to_package_dir, branch)
 
-    mock_update_llvm_next.assert_called_once_with(
-        abs_path_to_package, llvm_variant, git_hash, svn_version)
+    mock_update_llvm_next.assert_called_once_with(abs_path_to_package,
+                                                  llvm_variant, git_hash,
+                                                  svn_version)
 
     mock_uprev_symlink.assert_called_once_with(symlink_path_to_package)
 
+    mock_mask_contains.assert_called_once_with(chroot_path, git_hash)
+
     expected_commit_messages = [
         'llvm-next/tot: upgrade to a123testhash5 (r1000)\n',
         'The following packages have been updated:', 'path/to',
@@ -796,6 +829,48 @@ class UpdateLLVMHashTest(unittest.TestCase):
 
     mock_delete_repo.assert_called_once_with(path_to_package_dir, branch)
 
+  @mock.patch.object(subprocess, 'check_output', return_value=None)
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
+  def testEnsurePackageMaskContainsExisting(self, mock_llvm_version,
+                                            mock_git_add):
+    chroot_path = 'absolute/path/to/chroot'
+    git_hash = 'badf00d'
+    mock_llvm_version.return_value = '1234'
+    with mock.patch(
+        'update_chromeos_llvm_hash.open',
+        mock.mock_open(read_data='\n=sys-devel/llvm-1234.0_pre*\n'),
+        create=True) as mock_file:
+      update_chromeos_llvm_hash.EnsurePackageMaskContains(chroot_path, git_hash)
+      handle = mock_file()
+      handle.write.assert_not_called()
+    mock_llvm_version.assert_called_once_with(git_hash)
+
+    overlay_dir = 'absolute/path/to/chroot/src/third_party/chromiumos-overlay'
+    mask_path = overlay_dir + '/profiles/targets/chromeos/package.mask'
+    mock_git_add.assert_called_once_with(
+        ['git', '-C', overlay_dir, 'add', mask_path])
+
+  @mock.patch.object(subprocess, 'check_output', return_value=None)
+  @mock.patch.object(get_llvm_hash, 'GetLLVMMajorVersion')
+  def testEnsurePackageMaskContainsNotExisting(self, mock_llvm_version,
+                                               mock_git_add):
+    chroot_path = 'absolute/path/to/chroot'
+    git_hash = 'badf00d'
+    mock_llvm_version.return_value = '1234'
+    with mock.patch(
+        'update_chromeos_llvm_hash.open',
+        mock.mock_open(read_data='nothing relevant'),
+        create=True) as mock_file:
+      update_chromeos_llvm_hash.EnsurePackageMaskContains(chroot_path, git_hash)
+      handle = mock_file()
+      handle.write.assert_called_once_with('=sys-devel/llvm-1234.0_pre*\n')
+    mock_llvm_version.assert_called_once_with(git_hash)
+
+    overlay_dir = 'absolute/path/to/chroot/src/third_party/chromiumos-overlay'
+    mask_path = overlay_dir + '/profiles/targets/chromeos/package.mask'
+    mock_git_add.assert_called_once_with(
+        ['git', '-C', overlay_dir, 'add', mask_path])
+
 
 if __name__ == '__main__':
   unittest.main()