1 files changed, 44 insertions, 0 deletions
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()