rust_tools: Refactor to create and remove Rust versions independently

With this change, a user can run `./rust_uprev.py create` or
`./rust_uprev.py remove` to create a new Rust uprev or remove old
Rust version independently. Both supports using a state file to save
the progress.

BUG=chromium:1112551
TEST=unittest;generated a CL locally

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

2 files changed, 325 insertions, 261 deletions

diff --git a/rust_tools/rust_uprev.py b/rust_tools/rust_uprev.py
index f8c03331..50f85eb9 100755
--- a/rust_tools/rust_uprev.py
+++ b/rust_tools/rust_uprev.py
@@ -7,18 +7,25 @@
 """Tool to automatically generate a new Rust uprev CL.
 
 This tool is intended to automatically generate a CL to uprev Rust to a
-newer version in Chrome OS. It's based on
+newer version in Chrome OS, including creating a new Rust version or
+removing an old version. It's based on
 src/third_party/chromiumos-overlay/dev-lang/rust/UPGRADE.md. When using
 the tool, the progress can be saved to a JSON file, so the user can resume
-the process after a failing step is fixed. Example usage:
+the process after a failing step is fixed. Example usage to create a new
+version:
 
-1. (inside chroot) $ ./rust_tools/rust_uprev.py --rust_version 1.45.0 \
-                   --state_file /tmp/state-file.json
+1. (inside chroot) $ ./rust_tools/rust_uprev.py
+                     --state_file /tmp/state-file.json
+                     create --rust_version 1.45.0
 2. Step "compile rust" failed due to the patches can't apply to new version
 3. Manually fix the patches
 4. Execute the command in step 1 again.
 5. Iterate 1-4 for each failed step until the tool passes.
 
+Replace `create --rust_version 1.45.0` with `remove --rust_version 1.43.0`
+if you want to remove all 1.43.0 related stuff in the same CL. Remember to
+use a different state file if you choose to run different subcommands.
+
 See `--help` for all available options.
 """
 
@@ -36,7 +43,8 @@ import sys
 import tempfile
 from typing import Any, Callable, Dict, List, NamedTuple, Optional, T, Tuple
 
-from llvm_tools import chroot, git
+from llvm_tools import chroot
+RUST_PATH = '/mnt/host/source/src/third_party/chromiumos-overlay/dev-lang/rust'
 
 
 def get_command_output(command: List[str]) -> str:
@@ -77,6 +85,84 @@ class RustVersion(NamedTuple):
         int(m.group('major')), int(m.group('minor')), int(m.group('patch')))
 
 
+def parse_commandline_args() -> argparse.Namespace:
+  parser = argparse.ArgumentParser(
+      description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+  parser.add_argument(
+      '--state_file',
+      required=True,
+      help='A state file to hold previous completed steps. If the file '
+      'exists, it needs to be used together with --continue or --restart. '
+      'If not exist (do not use --continue in this case), we will create a '
+      'file for you.',
+  )
+  parser.add_argument(
+      '--restart',
+      action='store_true',
+      help='Restart from the first step. Ignore the completed steps in '
+      'the state file',
+  )
+  parser.add_argument(
+      '--continue',
+      dest='cont',
+      action='store_true',
+      help='Continue the steps from the state file',
+  )
+
+  subparsers = parser.add_subparsers(dest='subparser_name')
+  subparser_names = []
+
+  create_parser = subparsers.add_parser('create')
+  subparser_names.append('create')
+  create_parser.add_argument(
+      '--rust_version',
+      type=RustVersion.parse,
+      required=True,
+      help='Rust version to upgrade to, in the form a.b.c',
+  )
+  create_parser.add_argument(
+      '--template',
+      type=RustVersion.parse,
+      default=None,
+      help='A template to use for creating a Rust uprev from, in the form '
+      'a.b.c The ebuild has to exist in the chroot. If not specified, the '
+      'tool will use the current Rust version in the chroot as template.',
+  )
+  create_parser.add_argument(
+      '--skip_compile',
+      action='store_true',
+      help='Skip compiling rust to test the tool. Only for testing',
+  )
+
+  subparser_names.append('remove')
+  remove_parser = subparsers.add_parser('remove')
+  remove_parser.add_argument(
+      '--rust_version',
+      type=RustVersion.parse,
+      required=True,
+      help='Rust version to upgrade to, in the form a.b.c',
+  )
+
+  args = parser.parse_args()
+  if args.subparser_name not in subparser_names:
+    parser.error('one of %s must be specified' % subparser_names)
+
+  if args.cont and args.restart:
+    parser.error('Please select either --continue or --restart')
+
+  if os.path.exists(args.state_file):
+    if not args.cont and not args.restart:
+      parser.error('State file exists, so you should either --continue '
+                   'or --restart')
+  if args.cont and not os.path.exists(args.state_file):
+    parser.error('Indicate --continue but the state file does not exist')
+
+  if args.restart and os.path.exists(args.state_file):
+    os.remove(args.state_file)
+
+  return args
+
+
 def parse_stage0_file(new_version: RustVersion) -> Tuple[str, str, str]:
   # Find stage0 date, rustc and cargo
   stage0_file = get_command_output([
@@ -94,76 +180,55 @@ 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
-                           ) -> Tuple[str, RustVersion, RustVersion]:
-  a, b, c = json_input
-  return a, RustVersion(*b), RustVersion(*c)
+def prepare_uprev_from_json(json_input: Any) -> RustVersion:
+  return RustVersion(*json_input)
 
 
 def prepare_uprev(rust_version: RustVersion,
-                  reset: bool) -> Tuple[str, RustVersion, RustVersion]:
-  ebuild_path = get_command_output(['equery', 'w', 'rust'])
-  rust_path, ebuild_name = os.path.split(ebuild_path)
-  if reset:
-    subprocess.check_call(['git', 'reset', '--hard'], cwd=rust_path)
+                  template: Optional[RustVersion]) -> RustVersion:
+  if template is None:
     ebuild_path = get_command_output(['equery', 'w', 'rust'])
-    _, ebuild_name = os.path.split(ebuild_path)
-
-  current_version = RustVersion.parse(ebuild_name)
-  if rust_version <= current_version:
-    logging.info('Requested version %s is not newer than existing version %s.',
-                 rust_version, current_version)
-    return '', None, None
-
-  logging.info('Current Rust version is %s', current_version)
-  other_ebuilds = [
-      x for x in os.listdir(rust_path) if '.ebuild' in x and x != ebuild_name
-  ]
-  if len(other_ebuilds) != 1:
-    raise Exception('Expect exactly 1 previous version ebuild, '
-                    f'but actually found {other_ebuilds}')
-  # TODO(tcwang): Only support uprev from the older ebuild; need support to
-  # pick either version of the Rust to uprev from
-  old_version = RustVersion.parse(other_ebuilds[0])
-  # Prepare a repo branch for uprev
-  branch_name = f'rust-to-{rust_version}'
-  git.CreateBranch(rust_path, branch_name)
-  logging.info('Create a new repo branch %s', branch_name)
-  return rust_path, current_version, old_version
-
-
-def copy_patches(rust_path: str, old_version: RustVersion,
-                 current_version: RustVersion,
+    ebuild_name = os.path.basename(ebuild_path)
+    template_version = RustVersion.parse(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')
+    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
+
+
+def copy_patches(template_version: RustVersion,
                  new_version: RustVersion) -> None:
-  patch_path = os.path.join(rust_path, 'files')
+  patch_path = os.path.join(RUST_PATH, 'files')
   for f in os.listdir(patch_path):
-    if f'rust-{current_version}' not in f:
+    if f'rust-{template_version}' not in f:
       continue
     logging.info('Rename patch %s to new version', f)
-    new_name = f.replace(str(current_version), str(new_version))
+    new_name = f.replace(str(template_version), str(new_version))
     shutil.copyfile(
         os.path.join(patch_path, f),
         os.path.join(patch_path, new_name),
     )
 
   subprocess.check_call(['git', 'add', f'files/rust-{new_version}-*.patch'],
-                        cwd=rust_path)
-
-  subprocess.check_call(['git', 'rm', f'files/rust-{old_version}-*.patch'],
-                        cwd=rust_path)
+                        cwd=RUST_PATH)
 
 
-def rename_ebuild(rust_path: str, old_version: RustVersion,
-                  current_version: RustVersion,
+def create_ebuild(template_version: RustVersion,
                   new_version: RustVersion) -> str:
   shutil.copyfile(
-      os.path.join(rust_path, f'rust-{current_version}.ebuild'),
-      os.path.join(rust_path, f'rust-{new_version}.ebuild'))
+      os.path.join(RUST_PATH, f'rust-{template_version}.ebuild'),
+      os.path.join(RUST_PATH, f'rust-{new_version}.ebuild'))
   subprocess.check_call(['git', 'add', f'rust-{new_version}.ebuild'],
-                        cwd=rust_path)
-  subprocess.check_call(['git', 'rm', f'rust-{old_version}.ebuild'],
-                        cwd=rust_path)
-  return os.path.join(rust_path, f'rust-{new_version}.ebuild')
+                        cwd=RUST_PATH)
+  return os.path.join(RUST_PATH, f'rust-{new_version}.ebuild')
 
 
 def update_ebuild(ebuild_file: str, stage0_info: Tuple[str, str, str]) -> None:
@@ -235,38 +300,39 @@ def update_manifest(ebuild_file: str) -> None:
   flip_mirror_in_ebuild(ebuild_file, add=False)
 
 
-def upgrade_rust_packages(ebuild_file: str, old_version: RustVersion,
-                          current_version: RustVersion,
-                          new_version: RustVersion) -> None:
+def update_rust_packages(rust_version: RustVersion, add: bool) -> None:
   package_file = os.path.join(
-      os.path.dirname(ebuild_file),
-      '../../profiles/targets/chromeos/package.provided')
+      RUST_PATH, '../../profiles/targets/chromeos/package.provided')
   with open(package_file, encoding='utf-8') as f:
     contents = f.read()
-  old_str = f'dev-lang/rust-{old_version}'
-  current_str = f'dev-lang/rust-{current_version}'
-  new_str = f'dev-lang/rust-{new_version}'
-  if old_str not in contents or current_str not in contents:
-    raise Exception(f'Expect {old_str} and {current_str} to be in '
-                    'profiles/targets/chromeos/package.provided')
-  # Replace the two strings (old_str, current_str) with (current_str, new_str),
-  # so they are still ordered by rust versions
-  new_contents = contents.replace(current_str,
-                                  new_str).replace(old_str, current_str)
+  if add:
+    rust_packages_re = re.compile(r'dev-lang/rust-(\d+\.\d+\.\d+)')
+    rust_packages = rust_packages_re.findall(contents)
+    # Assume all the rust packages are in alphabetical order, so insert the new
+    # version to the place after the last rust_packages
+    new_str = f'dev-lang/rust-{rust_version}'
+    new_contents = contents.replace(rust_packages[-1],
+                                    f'{rust_packages[-1]}\n{new_str}')
+    logging.info('%s has been inserted into package.provided', new_str)
+  else:
+    old_str = f'dev-lang/rust-{rust_version}\n'
+    assert old_str in contents, f'{old_str!r} not found in package.provided'
+    new_contents = contents.replace(old_str, '')
+    logging.info('%s has been removed from package.provided', old_str)
+
   with open(package_file, 'w', encoding='utf-8') as f:
     f.write(new_contents)
-  logging.info('package.provided has been updated from %s, %s to %s, %s',
-               old_str, current_str, current_str, new_str)
 
 
-def update_virtual_rust(ebuild_file: str, old_version: RustVersion,
+def update_virtual_rust(template_version: RustVersion,
                         new_version: RustVersion) -> None:
-  virtual_rust_dir = os.path.join(
-      os.path.dirname(ebuild_file), '../../virtual/rust')
+  virtual_rust_dir = os.path.join(RUST_PATH, '../../virtual/rust')
   assert os.path.exists(virtual_rust_dir)
-  subprocess.check_call(
-      ['git', 'mv', f'rust-{old_version}.ebuild', f'rust-{new_version}.ebuild'],
-      cwd=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)
 
 
 def upload_to_localmirror(tempdir: str, rust_version: RustVersion) -> None:
@@ -310,64 +376,72 @@ def perform_step(state_file: pathlib.Path,
   return val
 
 
-def main():
-  if not chroot.InChroot():
-    raise RuntimeError('This script must be executed inside chroot')
+def create_rust_uprev(rust_version: RustVersion,
+                      template: Optional[RustVersion], skip_compile: bool,
+                      run_step: Callable[[
+                          str, Callable[[], T], Optional[Callable[[Any], T]],
+                          Optional[Callable[[T], Any]]
+                      ], T]) -> None:
+  stage0_info = run_step(
+      'parse stage0 file', lambda: parse_stage0_file(rust_version))
+  template_version = run_step(
+      'prepare uprev',
+      lambda: prepare_uprev(rust_version, template),
+      result_from_json=prepare_uprev_from_json,
+  )
+  if template_version is None:
+    return
 
-  logging.basicConfig(level=logging.INFO)
+  run_step('copy patches', lambda: copy_patches(template_version, rust_version))
+  ebuild_file = run_step(
+      'create ebuild', lambda: create_ebuild(template_version, 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(
+        tempdir, rust_version))
+  run_step('update manifest', lambda: update_manifest(ebuild_file))
+  if not skip_compile:
+    run_step('compile rust', lambda: rust_ebuild_command('compile'))
+    run_step('merge rust', lambda: rust_ebuild_command('merge', sudo=True))
+  run_step('insert version into rust packages', lambda: update_rust_packages(
+      rust_version, add=True))
+  run_step('upgrade virtual/rust', lambda: update_virtual_rust(
+      template_version, rust_version))
 
-  parser = argparse.ArgumentParser(
-      description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
-  parser.add_argument(
-      '--rust_version',
-      type=RustVersion.parse,
-      required=True,
-      help='Rust version to upgrade to, in the form a.b.c',
-  )
-  parser.add_argument(
-      '--state_file',
-      required=True,
-      help='A state file to hold previous completed steps. If the file '
-      'exists, it needs to be used together with --continue or --restart. '
-      'If not exist (do not use --continue in this case), we will create a '
-      'file for you.',
-  )
-  parser.add_argument(
-      '--skip_compile',
-      action='store_true',
-      help='Skip compiling rust to test the tool. Only for testing',
-  )
-  parser.add_argument(
-      '--restart',
-      action='store_true',
-      help='Restart from the first step. Ignore the completed steps in '
-      'the state file',
-  )
-  parser.add_argument(
-      '--continue',
-      dest='cont',
-      action='store_true',
-      help='Continue the steps from the state file',
-  )
 
-  args = parser.parse_args()
+def remove_files(filename: str, path: str) -> None:
+  subprocess.check_call(['git', 'rm', filename], cwd=path)
 
-  rust_version = args.rust_version
-  state_file = pathlib.Path(args.state_file)
-  tmp_state_file = pathlib.Path(args.state_file + '.tmp')
 
-  if args.cont and args.restart:
-    parser.error('Please select either --continue or --restart')
+def remove_rust_uprev(rust_version: RustVersion, run_step: Callable[[
+    str, Callable[[], T], Optional[Callable[[Any], T]], Optional[
+        Callable[[T], Any]]
+], T]) -> None:
+  run_step(
+      'remove patches', lambda: remove_files(
+          f'files/rust-{rust_version}-*.patch', RUST_PATH))
+  run_step('remove ebuild', lambda: remove_files(f'rust-{rust_version}.ebuild',
+                                                 RUST_PATH))
+  ebuild_file = get_command_output(['equery', 'w', 'rust'])
+  run_step('update manifest', lambda: update_manifest(ebuild_file))
+  run_step('remove version from rust packages', lambda: update_rust_packages(
+      rust_version, add=False))
+  run_step(
+      'remove virtual/rust', lambda: remove_files(
+          f'rust-{rust_version}.ebuild',
+          os.path.join(RUST_PATH, '../../virtual/rust')))
 
-  if os.path.exists(state_file):
-    if not args.cont and not args.restart:
-      parser.error('State file exists, so you should either --continue '
-                   'or --restart')
-  if args.cont and not os.path.exists(state_file):
-    parser.error('Indicate --continue but the state file does not exist')
 
-  if args.restart and os.path.exists(state_file):
-    os.remove(state_file)
+def main() -> None:
+  if not chroot.InChroot():
+    raise RuntimeError('This script must be executed inside chroot')
+
+  logging.basicConfig(level=logging.INFO)
+
+  args = parse_commandline_args()
+
+  state_file = pathlib.Path(args.state_file)
+  tmp_state_file = state_file.with_suffix('.tmp')
 
   try:
     with state_file.open(encoding='utf-8') as f:
@@ -384,38 +458,11 @@ def main():
     return perform_step(state_file, tmp_state_file, completed_steps, step_name,
                         step_fn, result_from_json, result_to_json)
 
-  stage0_info = run_step(
-      'parse stage0 file', lambda: parse_stage0_file(rust_version))
-  rust_path, current_version, old_version = run_step(
-      'prepare uprev',
-      lambda: prepare_uprev(rust_version, args.restart),
-      result_from_json=prepare_uprev_from_json,
-  )
-  if current_version is None:
-    return
-
-  current_version = RustVersion(*current_version)
-  old_version = RustVersion(*old_version)
-
-  run_step(
-      'copy patches', lambda: copy_patches(rust_path, old_version,
-                                           current_version, rust_version))
-  ebuild_file = run_step(
-      'rename ebuild', lambda: rename_ebuild(rust_path, old_version,
-                                             current_version, 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(
-        tempdir, rust_version))
-  run_step('update manifest', lambda: update_manifest(ebuild_file))
-  if not args.skip_compile:
-    run_step('compile rust', lambda: rust_ebuild_command('compile'))
-    run_step('merge rust', lambda: rust_ebuild_command('merge', sudo=True))
-  run_step(
-      'upgrade rust packages', lambda: upgrade_rust_packages(
-          ebuild_file, old_version, current_version, rust_version))
-  run_step('upgrade virtual/rust', lambda: update_virtual_rust(
-      ebuild_file, old_version, rust_version))
+  if args.subparser_name == 'create':
+    create_rust_uprev(args.rust_version, args.template, args.skip_compile,
+                      run_step)
+  else:
+    remove_rust_uprev(args.rust_version, run_step)
 
 
 if __name__ == '__main__':
diff --git a/rust_tools/rust_uprev_test.py b/rust_tools/rust_uprev_test.py
index 28b23bfa..e007b822 100755
--- a/rust_tools/rust_uprev_test.py
+++ b/rust_tools/rust_uprev_test.py
@@ -14,7 +14,6 @@ import unittest
 from unittest import mock
 
 import rust_uprev
-from llvm_tools import git
 
 
 class RustVersionTest(unittest.TestCase):
@@ -43,63 +42,65 @@ class RustVersionTest(unittest.TestCase):
 
 class PrepareUprevTest(unittest.TestCase):
   """Tests for prepare_uprev step in rust_uprev"""
-  mock_equery = '/path/to/rust/rust-1.2.3.ebuild'
-  mock_lsdir = ['rust-1.1.1.ebuild', 'rust-1.2.3.ebuild', 'an-unrelated-file']
 
-  @mock.patch.object(subprocess, 'check_call')
-  @mock.patch.object(git, 'CreateBranch')
+  def setUp(self):
+    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, 'get_command_output')
-  @mock.patch.object(os, 'listdir')
-  def test_success(self, mock_ls, mock_command, mock_git, mock_reset):
-    mock_ls.return_value = self.mock_lsdir
-    mock_command.return_value = self.mock_equery
-    input_version = rust_uprev.RustVersion(1, 3, 5)
-    expected = ('/path/to/rust', rust_uprev.RustVersion(1, 2, 3),
-                rust_uprev.RustVersion(1, 1, 1))
-    actual = rust_uprev.prepare_uprev(input_version, True)
+  def test_success_with_template(self, mock_command, mock_exists):
+    mock_exists.return_value = True
+    expected = self.version_old
+    actual = rust_uprev.prepare_uprev(
+        rust_version=self.version_new, template=self.version_old)
     self.assertEqual(expected, actual)
-    mock_reset.assert_called_once_with(['git', 'reset', '--hard'],
-                                       cwd='/path/to/rust')
-    mock_git.assert_called_once_with('/path/to/rust', 'rust-to-1.3.5')
-
-  @mock.patch.object(git, 'CreateBranch')
-  @mock.patch.object(
-      rust_uprev,
-      'get_command_output',
-      return_value='/path/to/rust/rust-1.2.3.ebuild')
-  @mock.patch.object(os, 'listdir')
-  def test_current_version_larger_failure(self, mock_ls, mock_command,
-                                          mock_git):
-    mock_command.return_value = self.mock_equery
-    input_version = rust_uprev.RustVersion(1, 1, 1)
-    rust_path, current, old = rust_uprev.prepare_uprev(input_version, False)
-    self.assertEqual(rust_path, '')
-    self.assertIsNone(current)
-    self.assertIsNone(old)
-    mock_ls.assert_not_called()
-    mock_git.assert_not_called()
-
-  @mock.patch.object(git, 'CreateBranch')
+    mock_command.assert_not_called()
+
+  @mock.patch.object(os.path, 'exists')
   @mock.patch.object(rust_uprev, 'get_command_output')
-  @mock.patch.object(os, 'listdir')
-  def test_more_than_two_ebuilds_fail(self, mock_ls, mock_command, mock_git):
-    mock_command.return_value = self.mock_equery
-    mock_ls.return_value = self.mock_lsdir + ['rust-1.0.0.ebuild']
-    input_version = rust_uprev.RustVersion(1, 3, 5)
-    with self.assertRaises(Exception) as context:
-      rust_uprev.prepare_uprev(input_version, False)
-    self.assertIn('Expect exactly 1 previous version ebuild',
-                  str(context.exception))
-    mock_git.assert_not_called()
+  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, 'get_command_output')
+  def test_return_none_with_template_larger_than_input(self, mock_command,
+                                                       mock_exists):
+    mock_exists.return_value = True
+    ret = rust_uprev.prepare_uprev(
+        rust_version=self.version_old, template=self.version_new)
+    self.assertIsNone(ret)
+    mock_command.assert_not_called()
+
+  @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
+    actual = rust_uprev.prepare_uprev(
+        rust_version=self.version_new, template=None)
+    self.assertEqual(expected, actual)
+    mock_command.assert_called_once_with(['equery', 'w', 'rust'])
+    mock_exists.assert_not_called()
+
+  @mock.patch.object(os.path, 'exists')
+  @mock.patch.object(rust_uprev, 'get_command_output')
+  def test_return_none_with_ebuild_larger_than_input(self, mock_command,
+                                                     mock_exists):
+    mock_command.return_value = f'/path/to/rust/rust-{self.version_new}.ebuild'
+    ret = rust_uprev.prepare_uprev(rust_version=self.version_old, template=None)
+    self.assertIsNone(ret)
+    mock_exists.assert_not_called()
 
   def test_prepare_uprev_from_json(self):
-    json_result = [
-        '/path/to/rust',
-        [1, 44, 0],
-        [1, 43, 0],
-    ]
-    expected = ('/path/to/rust', rust_uprev.RustVersion(1, 44, 0),
-                rust_uprev.RustVersion(1, 43, 0))
+    json_result = list(self.version_new)
+    expected = self.version_new
     actual = rust_uprev.prepare_uprev_from_json(json_result)
     self.assertEqual(expected, actual)
 
@@ -187,15 +188,49 @@ class UpdateManifestTest(unittest.TestCase):
          mock.call(ebuild_file, add=False)])
 
 
-class RustUprevOtherTests(unittest.TestCase):
+class UpdateRustPackagesTests(unittest.TestCase):
+  """Tests for update_rust_packages step."""
+
+  def setUp(self):
+    self.old_version = rust_uprev.RustVersion(1, 1, 0)
+    self.current_version = rust_uprev.RustVersion(1, 2, 3)
+    self.new_version = rust_uprev.RustVersion(1, 3, 5)
+    self.ebuild_file = os.path.join(rust_uprev.RUST_PATH,
+                                    'rust-{self.new_version}.ebuild')
+
+  def test_add_new_rust_packages(self):
+    package_before = (f'dev-lang/rust-{self.old_version}\n'
+                      f'dev-lang/rust-{self.current_version}')
+    package_after = (f'dev-lang/rust-{self.old_version}\n'
+                     f'dev-lang/rust-{self.current_version}\n'
+                     f'dev-lang/rust-{self.new_version}')
+    mock_open = mock.mock_open(read_data=package_before)
+    with mock.patch('builtins.open', mock_open):
+      rust_uprev.update_rust_packages(self.new_version, add=True)
+    mock_open.return_value.__enter__().write.assert_called_once_with(
+        package_after)
+
+  def test_remove_old_rust_packages(self):
+    package_before = (f'dev-lang/rust-{self.old_version}\n'
+                      f'dev-lang/rust-{self.current_version}\n'
+                      f'dev-lang/rust-{self.new_version}')
+    package_after = (f'dev-lang/rust-{self.current_version}\n'
+                     f'dev-lang/rust-{self.new_version}')
+    mock_open = mock.mock_open(read_data=package_before)
+    with mock.patch('builtins.open', mock_open):
+      rust_uprev.update_rust_packages(self.old_version, add=False)
+    mock_open.return_value.__enter__().write.assert_called_once_with(
+        package_after)
+
+
+class RustUprevOtherStagesTests(unittest.TestCase):
   """Tests for other steps in rust_uprev"""
 
   def setUp(self):
-    self.rust_path = '/path/to/rust'
     self.old_version = rust_uprev.RustVersion(1, 1, 0)
     self.current_version = rust_uprev.RustVersion(1, 2, 3)
     self.new_version = rust_uprev.RustVersion(1, 3, 5)
-    self.ebuild_file = os.path.join(self.rust_path,
+    self.ebuild_file = os.path.join(rust_uprev.RUST_PATH,
                                     'rust-{self.new_version}.ebuild')
 
   @mock.patch.object(rust_uprev, 'get_command_output')
@@ -227,66 +262,48 @@ class RustUprevOtherTests(unittest.TestCase):
         f'rust-{self.current_version}-patch-1.patch',
         f'rust-{self.current_version}-patch-2-new.patch'
     ]
-    rust_uprev.copy_patches(self.rust_path, self.old_version,
-                            self.current_version, self.new_version)
+    rust_uprev.copy_patches(self.current_version, self.new_version)
     mock_copy.assert_has_calls([
         mock.call(
-            os.path.join(self.rust_path, 'files',
+            os.path.join(rust_uprev.RUST_PATH, 'files',
                          f'rust-{self.current_version}-patch-1.patch'),
-            os.path.join(self.rust_path, 'files',
+            os.path.join(rust_uprev.RUST_PATH, 'files',
                          f'rust-{self.new_version}-patch-1.patch'),
         ),
         mock.call(
-            os.path.join(self.rust_path, 'files',
+            os.path.join(rust_uprev.RUST_PATH, 'files',
                          f'rust-{self.current_version}-patch-2-new.patch'),
-            os.path.join(self.rust_path, 'files',
+            os.path.join(rust_uprev.RUST_PATH, 'files',
                          f'rust-{self.new_version}-patch-2-new.patch'))
     ])
-    mock_call.assert_has_calls([
-        mock.call(['git', 'add', f'files/rust-{self.new_version}-*.patch'],
-                  cwd=self.rust_path),
-        mock.call(['git', 'rm', f'files/rust-{self.old_version}-*.patch'],
-                  cwd=self.rust_path)
-    ])
+    mock_call.assert_called_once_with(
+        ['git', 'add', f'files/rust-{self.new_version}-*.patch'],
+        cwd=rust_uprev.RUST_PATH)
 
   @mock.patch.object(shutil, 'copyfile')
   @mock.patch.object(subprocess, 'check_call')
-  def test_rename_ebuild(self, mock_call, mock_copy):
-    rust_uprev.rename_ebuild(self.rust_path, self.old_version,
-                             self.current_version, self.new_version)
+  def test_create_ebuild(self, mock_call, mock_copy):
+    rust_uprev.create_ebuild(self.current_version, self.new_version)
     mock_copy.assert_called_once_with(
-        os.path.join(self.rust_path, f'rust-{self.current_version}.ebuild'),
-        os.path.join(self.rust_path, f'rust-{self.new_version}.ebuild'))
-    mock_call.assert_has_calls([
-        mock.call(['git', 'add', f'rust-{self.new_version}.ebuild'],
-                  cwd=self.rust_path),
-        mock.call(['git', 'rm', f'rust-{self.old_version}.ebuild'],
-                  cwd=self.rust_path)
-    ])
-
-  def test_upgrade_rust_packages(self):
-    package_before = (f'dev-lang/rust-{self.old_version}\n'
-                      f'dev-lang/rust-{self.current_version}')
-    package_after = (f'dev-lang/rust-{self.current_version}\n'
-                     f'dev-lang/rust-{self.new_version}')
-    mock_open = mock.mock_open(read_data=package_before)
-    with mock.patch('builtins.open', mock_open):
-      rust_uprev.upgrade_rust_packages(self.ebuild_file, self.old_version,
-                                       self.current_version, self.new_version)
-    mock_open.return_value.__enter__().write.assert_called_once_with(
-        package_after)
+        os.path.join(rust_uprev.RUST_PATH,
+                     f'rust-{self.current_version}.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'],
+        cwd=rust_uprev.RUST_PATH)
 
   @mock.patch.object(os.path, 'exists', return_value=True)
+  @mock.patch.object(shutil, 'copyfile')
   @mock.patch.object(subprocess, 'check_call')
-  def test_update_virtual_rust(self, mock_call, _):
-    rust_uprev.update_virtual_rust(self.ebuild_file, self.old_version,
-                                   self.new_version)
-    mock_call.assert_called_once_with([
-        'git', 'mv', f'rust-{self.old_version}.ebuild',
-        f'rust-{self.new_version}.ebuild'
-    ],
-                                      cwd=os.path.join(self.rust_path,
-                                                       '../../virtual/rust'))
+  def test_update_virtual_rust(self, mock_call, mock_copy, mock_exists):
+    virtual_rust_dir = os.path.join(rust_uprev.RUST_PATH, '../../virtual/rust')
+    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(subprocess, 'check_call')
   def test_upload_to_localmirror(self, mock_call):