1 files changed, 81 insertions, 25 deletions
diff --git a/llvm_tools/get_llvm_hash_unittest.py b/llvm_tools/get_llvm_hash_unittest.py
index 66685ca1..fda002ce 100755
--- a/llvm_tools/get_llvm_hash_unittest.py
+++ b/llvm_tools/get_llvm_hash_unittest.py
@@ -6,11 +6,16 @@
 """Unit tests for retrieving the LLVM hash."""
 
 import contextlib
+from pathlib import Path
+import shutil
 import subprocess
+import tempfile
+import textwrap
 import unittest
 from unittest import mock
 
 import get_llvm_hash
+import llvm_next
 import subprocess_helpers
 
 
@@ -28,6 +33,15 @@ def mock_run_results(returncode: int, stderr: bytes) -> mock.MagicMock:
 class TestGetLLVMHash(unittest.TestCase):
     """The LLVMHash test class."""
 
+    def make_tempdir(self):
+        d = Path(tempfile.mkdtemp(prefix="get_llvm_hash_unittest_"))
+        self.addCleanup(shutil.rmtree, d)
+        return d
+
+    def setUp(self):
+        # We mock out quite a bit. Ensure every test is self-contained.
+        get_llvm_hash.GetLLVMMajorVersion.cache_clear()
+
     @mock.patch.object(subprocess, "run")
     def testCloneRepoSucceedsWhenGitSucceeds(self, run_mock):
         run_mock.return_value = mock_run_results(returncode=0, stderr=b"")
@@ -121,37 +135,79 @@ class TestGetLLVMHash(unittest.TestCase):
 
     def testParseLLVMMajorVersionInvalid(self):
         invalid_cmakelist = "invalid cmakelist.txt contents"
-        with self.assertRaises(ValueError):
+        self.assertIsNone(
             get_llvm_hash.ParseLLVMMajorVersion(invalid_cmakelist)
+        )
+
+    @mock.patch.object(get_llvm_hash, "GetAndUpdateLLVMProjectInLLVMTools")
+    @mock.patch.object(subprocess_helpers, "CheckCommand")
+    def testGetLLVMMajorVersionWithOldPath(
+        self,
+        _mock_check_command,
+        mock_update_project,
+    ):
+        src_dir = self.make_tempdir()
+        mock_update_project.return_value = str(src_dir)
+
+        cmakelists = Path(src_dir) / "llvm" / "CMakeLists.txt"
+        cmakelists.parent.mkdir(parents=True)
+        cmakelists.write_text(
+            textwrap.dedent(
+                """
+                if(NOT DEFINED LLVM_VERSION_MAJOR)
+                  set(LLVM_VERSION_MAJOR 12345)
+                endif()
+                """
+            ),
+            encoding="utf-8",
+        )
+        self.assertEqual(get_llvm_hash.GetLLVMMajorVersion(), "12345")
 
     @mock.patch.object(get_llvm_hash, "GetAndUpdateLLVMProjectInLLVMTools")
-    @mock.patch.object(get_llvm_hash, "ParseLLVMMajorVersion")
     @mock.patch.object(subprocess_helpers, "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(
+    def testGetLLVMMajorVersionWithNewPath(
         self,
-        mock_checkout_branch,
-        mock_git_checkout,
-        mock_major_version,
-        mock_llvm_project_path,
+        _mock_check_command,
+        mock_update_project,
     ):
-        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"
+        src_dir = self.make_tempdir()
+        mock_update_project.return_value = str(src_dir)
+
+        old_cmakelists = Path(src_dir) / "llvm" / "CMakeLists.txt"
+        old_cmakelists.parent.mkdir(parents=True)
+        old_cmakelists.write_text(
+            textwrap.dedent(
+                """
+                Some text
+                that has
+                nothing to do with
+                LLVM_VERSION_MAJOR
+                """
+            ),
+            encoding="utf-8",
+        )
+
+        new_cmakelists = (
+            Path(src_dir) / "cmake" / "Modules" / "LLVMVersion.cmake"
+        )
+        new_cmakelists.parent.mkdir(parents=True)
+        new_cmakelists.write_text(
+            textwrap.dedent(
+                """
+                if(NOT DEFINED LLVM_VERSION_MAJOR)
+                  set(LLVM_VERSION_MAJOR 5432)
+                endif()
+                """
+            ),
+            encoding="utf-8",
+        )
+
+        self.assertEqual(get_llvm_hash.GetLLVMMajorVersion(), "5432")
+
+    def testGetLLVMNextHash(self):
+        self.assertEqual(
+            get_llvm_hash.LLVMHash().GetCrOSLLVMNextHash(),
+            llvm_next.LLVM_NEXT_HASH,
         )