rust_uprev: accept -r* in virtual/rust version numbers

The ebuilds in virtual/rust can have patch (-r123) version numbers,
which currently cause rust_uprev to fail. This CL updates rust_uprev
to handle those.

BUG=chromium:1159066
TEST=python3 ./rust_tools/rust_uprev_test.py

Change-Id: I76ccaf553ea192db57e17ff3df9b02eee4105010
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2593514
Commit-Queue: Bob Haarman <inglorion@chromium.org>
Tested-by: Bob Haarman <inglorion@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>
Reviewed-by: Tiancong Wang <tcwang@google.com>

2 files changed, 73 insertions, 14 deletions

diff --git a/rust_tools/rust_uprev.py b/rust_tools/rust_uprev.py
index 3c0ad012..b4020ea1 100755
--- a/rust_tools/rust_uprev.py
+++ b/rust_tools/rust_uprev.py
@@ -36,6 +36,7 @@ See `--help` for all available options.
 # pylint: disable=cros-logging-import
 
 import argparse
+import glob
 import pathlib
 import json
 import logging
@@ -91,6 +92,20 @@ class RustVersion(NamedTuple):
         int(m.group('major')), int(m.group('minor')), int(m.group('patch')))
 
 
+def find_virtual_rust_ebuild(version: RustVersion) -> str:
+  """Finds the virtual/rust ebuild for a given RustVersion.
+
+  This finds 1.2.3 and also 1.2.3-r4. It expects that there will
+  be exactly one match, and will assert if that is not the case.
+  """
+  virtual_rust_dir = os.path.join(RUST_PATH, '../../virtual/rust')
+  pattern = os.path.join(virtual_rust_dir, f'rust-{version}*.ebuild')
+  matches = glob.glob(pattern)
+  # We expect exactly one match.
+  assert len(matches) == 1, matches
+  return matches[0]
+
+
 def parse_commandline_args() -> argparse.Namespace:
   parser = argparse.ArgumentParser(
       description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
@@ -367,13 +382,12 @@ def update_rust_packages(rust_version: RustVersion, add: bool) -> None:
 
 def update_virtual_rust(template_version: RustVersion,
                         new_version: RustVersion) -> None:
-  virtual_rust_dir = os.path.join(RUST_PATH, '../../virtual/rust')
-  assert os.path.exists(virtual_rust_dir)
-  shutil.copyfile(
-      os.path.join(virtual_rust_dir, f'rust-{template_version}.ebuild'),
-      os.path.join(virtual_rust_dir, f'rust-{new_version}.ebuild'))
-  subprocess.check_call(['git', 'add', f'rust-{new_version}.ebuild'],
-                        cwd=virtual_rust_dir)
+  template_ebuild = find_virtual_rust_ebuild(template_version)
+  virtual_rust_dir = os.path.dirname(template_ebuild)
+  new_name = f'rust-{new_version}.ebuild'
+  new_ebuild = os.path.join(virtual_rust_dir, new_name)
+  shutil.copyfile(template_ebuild, new_ebuild)
+  subprocess.check_call(['git', 'add', new_name], cwd=virtual_rust_dir)
 
 
 def upload_single_tarball(rust_url: str, tarfile_name: str,
@@ -569,10 +583,12 @@ def remove_rust_uprev(rust_version: Optional[RustVersion],
       ebuild_file))
   run_step('remove version from rust packages', lambda: update_rust_packages(
       delete_version, add=False))
-  run_step(
-      'remove virtual/rust', lambda: remove_files(
-          f'rust-{delete_version}.ebuild',
-          os.path.join(RUST_PATH, '../../virtual/rust')))
+  run_step('remove virtual/rust', lambda: remove_virtual_rust(delete_version))
+
+
+def remove_virtual_rust(delete_version: RustVersion) -> None:
+  dirname, basename = os.path.split(find_virtual_rust_ebuild(delete_version))
+  subprocess.check_call(['git', 'rm', basename], cwd=dirname)
 
 
 def create_new_repo(rust_version: RustVersion) -> None:
diff --git a/rust_tools/rust_uprev_test.py b/rust_tools/rust_uprev_test.py
index 0de785ea..acf700b8 100755
--- a/rust_tools/rust_uprev_test.py
+++ b/rust_tools/rust_uprev_test.py
@@ -7,6 +7,7 @@
 """Tests for rust_uprev.py"""
 
 # pylint: disable=cros-logging-import
+import glob
 import os
 import shutil
 import subprocess
@@ -389,18 +390,60 @@ class RustUprevOtherStagesTests(unittest.TestCase):
         ['git', 'add', f'rust-{self.new_version}.ebuild'],
         cwd=rust_uprev.RUST_PATH)
 
-  @mock.patch.object(os.path, 'exists', return_value=True)
+  @mock.patch.object(glob, 'glob')
+  @mock.patch.object(subprocess, 'check_call')
+  def test_remove_virtual_rust(self, mock_call, mock_glob):
+    virtual_rust_dir = os.path.join(rust_uprev.RUST_PATH, '../../virtual/rust')
+    mock_glob.return_value = [
+        os.path.join(virtual_rust_dir, f'rust-{self.old_version}.ebuild'),
+    ]
+    rust_uprev.remove_virtual_rust(self.old_version)
+    mock_call.assert_called_once_with(
+        ['git', 'rm', f'rust-{self.old_version}.ebuild'], cwd=virtual_rust_dir)
+
+  @mock.patch.object(glob, 'glob')
+  @mock.patch.object(subprocess, 'check_call')
+  def test_remove_virtual_rust_patch(self, mock_call, mock_glob):
+    virtual_rust_dir = os.path.join(rust_uprev.RUST_PATH, '../../virtual/rust')
+    mock_glob.return_value = [
+        os.path.join(virtual_rust_dir, f'rust-{self.old_version}-r3.ebuild'),
+    ]
+    rust_uprev.remove_virtual_rust(self.old_version)
+    mock_call.assert_called_once_with(
+        ['git', 'rm', f'rust-{self.old_version}-r3.ebuild'],
+        cwd=virtual_rust_dir)
+
+  @mock.patch.object(glob, 'glob')
   @mock.patch.object(shutil, 'copyfile')
   @mock.patch.object(subprocess, 'check_call')
-  def test_update_virtual_rust(self, mock_call, mock_copy, mock_exists):
+  def test_update_virtual_rust(self, mock_call, mock_copy, mock_glob):
     virtual_rust_dir = os.path.join(rust_uprev.RUST_PATH, '../../virtual/rust')
+    mock_glob.return_value = [
+        os.path.join(virtual_rust_dir, f'rust-{self.current_version}.ebuild'),
+    ]
     rust_uprev.update_virtual_rust(self.current_version, self.new_version)
     mock_call.assert_called_once_with(
         ['git', 'add', f'rust-{self.new_version}.ebuild'], cwd=virtual_rust_dir)
     mock_copy.assert_called_once_with(
         os.path.join(virtual_rust_dir, f'rust-{self.current_version}.ebuild'),
         os.path.join(virtual_rust_dir, f'rust-{self.new_version}.ebuild'))
-    mock_exists.assert_called_once_with(virtual_rust_dir)
+
+  @mock.patch.object(glob, 'glob')
+  @mock.patch.object(shutil, 'copyfile')
+  @mock.patch.object(subprocess, 'check_call')
+  def test_update_virtual_rust_patched(self, mock_call, mock_copy, mock_glob):
+    virtual_rust_dir = os.path.join(rust_uprev.RUST_PATH, '../../virtual/rust')
+    mock_glob.return_value = [
+        os.path.join(virtual_rust_dir,
+                     f'rust-{self.current_version}-r3.ebuild'),
+    ]
+    rust_uprev.update_virtual_rust(self.current_version, self.new_version)
+    mock_call.assert_called_once_with(
+        ['git', 'add', f'rust-{self.new_version}.ebuild'], cwd=virtual_rust_dir)
+    mock_copy.assert_called_once_with(
+        os.path.join(virtual_rust_dir,
+                     f'rust-{self.current_version}-r3.ebuild'),
+        os.path.join(virtual_rust_dir, f'rust-{self.new_version}.ebuild'))
 
   @mock.patch.object(os, 'listdir')
   def test_find_oldest_rust_version_in_chroot_pass(self, mock_ls):