rust_uprev: support -rN ebuilds

rust_uprev currently doesn't fully support ebuilds with -r123 after
Rust's version number. This CL aims to help us support those better.

BUG=chromium:1136579, chromium:1112551
TEST=Ran the script

Change-Id: I22a09ec3f5d074be547783e4cb9b702d4aad0bd2
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2462925
Reviewed-by: Tiancong Wang <tcwang@google.com>
Tested-by: George Burgess <gbiv@chromium.org>

2 files changed, 85 insertions, 62 deletions

diff --git a/rust_tools/rust_uprev.py b/rust_tools/rust_uprev.py
index db907149..b4fcaf18 100755
--- a/rust_tools/rust_uprev.py
+++ b/rust_tools/rust_uprev.py
@@ -71,6 +71,7 @@ class RustVersion(NamedTuple):
                           r'(?P<major>\d+)\.'
                           r'(?P<minor>\d+)\.'
                           r'(?P<patch>\d+)'
+                          r'(:?-r\d+)?'
                           r'\.ebuild$')
     m = input_re.match(ebuild_name)
     assert m, f'failed to parse {ebuild_name!r}'
@@ -226,28 +227,25 @@ def parse_stage0_file(new_version: RustVersion) -> Tuple[str, str, str]:
   return stage0_date, stage0_rustc, stage0_cargo
 
 
-def prepare_uprev_from_json(json_input: Any) -> RustVersion:
-  return RustVersion(*json_input)
-
-
-def prepare_uprev(rust_version: RustVersion,
-                  template: Optional[RustVersion]) -> RustVersion:
+def prepare_uprev(rust_version: RustVersion, template: Optional[RustVersion]
+                 ) -> Optional[Tuple[RustVersion, str]]:
   if template is None:
     ebuild_path = get_command_output(['equery', 'w', 'rust'])
     ebuild_name = os.path.basename(ebuild_path)
-    template_version = RustVersion.parse(ebuild_name)
+    template_version = RustVersion.parse_from_ebuild(ebuild_name)
   else:
-    if not os.path.exists(os.path.join(RUST_PATH, f'rust-{template}.ebuild')):
-      raise ValueError(f'Template ebuild file {template} does not exist')
+    ebuild_path = find_ebuild_for_rust_version(template)
     template_version = template
+
   if rust_version <= template_version:
     logging.info(
         'Requested version %s is not newer than the template version %s.',
         rust_version, template_version)
     return None
 
-  logging.info('Template Rust version is %s', template_version)
-  return template_version
+  logging.info('Template Rust version is %s (ebuild: %r)', template_version,
+               ebuild_path)
+  return template_version, ebuild_path
 
 
 def copy_patches(template_version: RustVersion,
@@ -267,11 +265,9 @@ def copy_patches(template_version: RustVersion,
                         cwd=RUST_PATH)
 
 
-def create_ebuild(template_version: RustVersion,
-                  new_version: RustVersion) -> str:
-  shutil.copyfile(
-      os.path.join(RUST_PATH, f'rust-{template_version}.ebuild'),
-      os.path.join(RUST_PATH, f'rust-{new_version}.ebuild'))
+def create_ebuild(template_ebuild: str, new_version: RustVersion) -> str:
+  shutil.copyfile(template_ebuild,
+                  os.path.join(RUST_PATH, f'rust-{new_version}.ebuild'))
   subprocess.check_call(['git', 'add', f'rust-{new_version}.ebuild'],
                         cwd=RUST_PATH)
   return os.path.join(RUST_PATH, f'rust-{new_version}.ebuild')
@@ -473,22 +469,29 @@ def perform_step(state_file: pathlib.Path,
   return val
 
 
+def prepare_uprev_from_json(obj: Any) -> Optional[Tuple[RustVersion, str]]:
+  if not obj:
+    return None
+  version, ebuild_path = obj
+  return RustVersion(*version), ebuild_path
+
+
 def create_rust_uprev(rust_version: RustVersion,
-                      template: Optional[RustVersion], skip_compile: bool,
-                      run_step: Callable[[], T]) -> None:
+                      maybe_template_version: Optional[RustVersion],
+                      skip_compile: bool, run_step: Callable[[], T]) -> None:
   stage0_info = run_step(
       'parse stage0 file', lambda: parse_stage0_file(rust_version))
-  template_version = run_step(
+  template_version, template_ebuild = run_step(
       'prepare uprev',
-      lambda: prepare_uprev(rust_version, template),
+      lambda: prepare_uprev(rust_version, maybe_template_version),
       result_from_json=prepare_uprev_from_json,
   )
-  if template_version is None:
+  if template_ebuild is None:
     return
 
   run_step('copy patches', lambda: copy_patches(template_version, rust_version))
   ebuild_file = run_step(
-      'create ebuild', lambda: create_ebuild(template_version, rust_version))
+      'create ebuild', lambda: create_ebuild(template_ebuild, rust_version))
   run_step('update ebuild', lambda: update_ebuild(ebuild_file, stage0_info))
   with tempfile.TemporaryDirectory(dir='/tmp') as tempdir:
     run_step('upload_to_localmirror', lambda: upload_to_localmirror(
@@ -504,32 +507,52 @@ def create_rust_uprev(rust_version: RustVersion,
       template_version, rust_version))
 
 
-def find_oldest_rust_version_inchroot() -> RustVersion:
-  rust_versions = [
-      RustVersion.parse(x) for x in os.listdir(RUST_PATH) if '.ebuild' in x
-  ]
+def find_rust_versions_in_chroot() -> List[Tuple[RustVersion, str]]:
+  return [(RustVersion.parse_from_ebuild(x), os.path.join(RUST_PATH, x))
+          for x in os.listdir(RUST_PATH)
+          if x.endswith('.ebuild')]
+
 
+def find_oldest_rust_version_in_chroot() -> Tuple[RustVersion, str]:
+  rust_versions = find_rust_versions_in_chroot()
   if len(rust_versions) <= 1:
     raise RuntimeError('Expect to find more than one Rust versions')
   return min(rust_versions)
 
 
+def find_ebuild_for_rust_version(version: RustVersion) -> str:
+  rust_ebuilds = [
+      ebuild for x, ebuild in find_rust_versions_in_chroot() if x == version
+  ]
+  if not rust_ebuilds:
+    raise ValueError(f'No Rust ebuilds found matching {version}')
+  if len(rust_ebuilds) > 1:
+    raise ValueError(f'Multiple Rust ebuilds found matching {version}: '
+                     f'{rust_ebuilds}')
+  return rust_ebuilds[0]
+
+
 def remove_files(filename: str, path: str) -> None:
   subprocess.check_call(['git', 'rm', filename], cwd=path)
 
 
 def remove_rust_uprev(rust_version: Optional[RustVersion],
                       run_step: Callable[[], T]) -> None:
-  delete_version = run_step(
+
+  def find_desired_rust_version():
+    if rust_version:
+      return rust_version, find_ebuild_for_rust_version(rust_version)
+    return find_oldest_rust_version_in_chroot()
+
+  delete_version, delete_ebuild = run_step(
       'find rust version to delete',
-      lambda: rust_version or find_oldest_rust_version_inchroot(),
+      find_desired_rust_version,
       result_from_json=prepare_uprev_from_json,
   )
   run_step(
       'remove patches', lambda: remove_files(
           f'files/rust-{delete_version}-*.patch', RUST_PATH))
-  run_step('remove ebuild', lambda: remove_files(
-      f'rust-{delete_version}.ebuild', RUST_PATH))
+  run_step('remove ebuild', lambda: remove_files(delete_ebuild, RUST_PATH))
   ebuild_file = get_command_output(['equery', 'w', 'rust'])
   run_step('update manifest to delete old version', lambda: update_manifest(
       ebuild_file))
diff --git a/rust_tools/rust_uprev_test.py b/rust_tools/rust_uprev_test.py
index c2f071f2..d7efdce2 100755
--- a/rust_tools/rust_uprev_test.py
+++ b/rust_tools/rust_uprev_test.py
@@ -35,6 +35,9 @@ class RustVersionTest(unittest.TestCase):
     actual = rust_uprev.RustVersion.parse_from_ebuild('rust-1.2.3.ebuild')
     self.assertEqual(expected, actual)
 
+    actual = rust_uprev.RustVersion.parse_from_ebuild('rust-1.2.3-r1.ebuild')
+    self.assertEqual(expected, actual)
+
   def test_parse_fail(self):
     with self.assertRaises(AssertionError) as context:
       rust_uprev.RustVersion.parse('invalid-rust-1.2.3')
@@ -49,32 +52,25 @@ class PrepareUprevTest(unittest.TestCase):
     self.version_old = rust_uprev.RustVersion(1, 2, 3)
     self.version_new = rust_uprev.RustVersion(1, 3, 5)
 
-  @mock.patch.object(os.path, 'exists')
+  @mock.patch.object(
+      rust_uprev,
+      'find_ebuild_for_rust_version',
+      return_value='/path/to/ebuild')
   @mock.patch.object(rust_uprev, 'get_command_output')
-  def test_success_with_template(self, mock_command, mock_exists):
-    mock_exists.return_value = True
-    expected = self.version_old
+  def test_success_with_template(self, mock_command, mock_find_ebuild):
+    expected = (self.version_old, '/path/to/ebuild')
     actual = rust_uprev.prepare_uprev(
         rust_version=self.version_new, template=self.version_old)
     self.assertEqual(expected, actual)
     mock_command.assert_not_called()
 
-  @mock.patch.object(os.path, 'exists')
-  @mock.patch.object(rust_uprev, 'get_command_output')
-  def test_fail_with_template_not_exist(self, mock_command, mock_exists):
-    mock_exists.return_value = False
-    with self.assertRaises(ValueError) as context:
-      rust_uprev.prepare_uprev(
-          rust_version=self.version_new, template=self.version_old)
-    self.assertEqual(f'Template ebuild file {self.version_old} does not exist',
-                     str(context.exception))
-    mock_command.assert_not_called()
-
-  @mock.patch.object(os.path, 'exists')
+  @mock.patch.object(
+      rust_uprev,
+      'find_ebuild_for_rust_version',
+      return_value='/path/to/ebuild')
   @mock.patch.object(rust_uprev, 'get_command_output')
   def test_return_none_with_template_larger_than_input(self, mock_command,
-                                                       mock_exists):
-    mock_exists.return_value = True
+                                                       _mock_find_ebuild):
     ret = rust_uprev.prepare_uprev(
         rust_version=self.version_old, template=self.version_new)
     self.assertIsNone(ret)
@@ -83,8 +79,9 @@ class PrepareUprevTest(unittest.TestCase):
   @mock.patch.object(os.path, 'exists')
   @mock.patch.object(rust_uprev, 'get_command_output')
   def test_success_without_template(self, mock_command, mock_exists):
-    mock_command.return_value = f'/path/to/rust/rust-{self.version_old}.ebuild'
-    expected = self.version_old
+    rust_ebuild_path = f'/path/to/rust/rust-{self.version_old}-r3.ebuild'
+    mock_command.return_value = rust_ebuild_path
+    expected = (self.version_old, rust_ebuild_path)
     actual = rust_uprev.prepare_uprev(
         rust_version=self.version_new, template=None)
     self.assertEqual(expected, actual)
@@ -101,8 +98,9 @@ class PrepareUprevTest(unittest.TestCase):
     mock_exists.assert_not_called()
 
   def test_prepare_uprev_from_json(self):
-    json_result = list(self.version_new)
-    expected = self.version_new
+    ebuild_path = '/path/to/the/ebuild'
+    json_result = (list(self.version_new), ebuild_path)
+    expected = (self.version_new, ebuild_path)
     actual = rust_uprev.prepare_uprev_from_json(json_result)
     self.assertEqual(expected, actual)
 
@@ -368,10 +366,10 @@ class RustUprevOtherStagesTests(unittest.TestCase):
   @mock.patch.object(shutil, 'copyfile')
   @mock.patch.object(subprocess, 'check_call')
   def test_create_ebuild(self, mock_call, mock_copy):
-    rust_uprev.create_ebuild(self.current_version, self.new_version)
+    template_ebuild = f'/path/to/rust-{self.current_version}-r2.ebuild'
+    rust_uprev.create_ebuild(template_ebuild, self.new_version)
     mock_copy.assert_called_once_with(
-        os.path.join(rust_uprev.RUST_PATH,
-                     f'rust-{self.current_version}.ebuild'),
+        template_ebuild,
         os.path.join(rust_uprev.RUST_PATH, f'rust-{self.new_version}.ebuild'))
     mock_call.assert_called_once_with(
         ['git', 'add', f'rust-{self.new_version}.ebuild'],
@@ -391,21 +389,23 @@ class RustUprevOtherStagesTests(unittest.TestCase):
     mock_exists.assert_called_once_with(virtual_rust_dir)
 
   @mock.patch.object(os, 'listdir')
-  def test_find_oldest_rust_version_inchroot_pass(self, mock_ls):
+  def test_find_oldest_rust_version_in_chroot_pass(self, mock_ls):
+    oldest_version_name = f'rust-{self.old_version}.ebuild'
     mock_ls.return_value = [
-        f'rust-{self.old_version}.ebuild',
-        f'rust-{self.current_version}.ebuild', f'rust-{self.new_version}.ebuild'
+        oldest_version_name, f'rust-{self.current_version}.ebuild',
+        f'rust-{self.new_version}.ebuild'
     ]
-    actual = rust_uprev.find_oldest_rust_version_inchroot()
-    expected = self.old_version
+    actual = rust_uprev.find_oldest_rust_version_in_chroot()
+    expected = (self.old_version,
+                os.path.join(rust_uprev.RUST_PATH, oldest_version_name))
     self.assertEqual(expected, actual)
 
   @mock.patch.object(os, 'listdir')
-  def test_find_oldest_rust_version_inchroot_fail_with_only_one_ebuild(
+  def test_find_oldest_rust_version_in_chroot_fail_with_only_one_ebuild(
       self, mock_ls):
     mock_ls.return_value = [f'rust-{self.new_version}.ebuild']
     with self.assertRaises(RuntimeError) as context:
-      rust_uprev.find_oldest_rust_version_inchroot()
+      rust_uprev.find_oldest_rust_version_in_chroot()
     self.assertEqual('Expect to find more than one Rust versions',
                      str(context.exception))