rust_tools: Add a tool to automatically generate a Rust uprev

This is the first pass to add a tool to convert rust upgrade
process into a program that can generate CLs autoamatically.
The tool is resumable, so when patches fail to apply, the user
can fix the patch, resume the tool until no problems exist.

The current tool only supports uprev from the earlier Rust ebuild
in chroot and will need future passes to add more features.

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

Change-Id: I77fccd14c69548824a8b235d756357aee0c42ef2
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2339593
Commit-Queue: Tiancong Wang <tcwang@google.com>
Tested-by: Tiancong Wang <tcwang@google.com>
Reviewed-by: Bob Haarman <inglorion@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>
diff --git a/rust_tools/rust_uprev.py b/rust_tools/rust_uprev.py
new file mode 100755
index 0000000..f8c0333
--- /dev/null
+++ b/rust_tools/rust_uprev.py
@@ -0,0 +1,422 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright 2020 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""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
+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:
+
+1. (inside chroot) $ ./rust_tools/rust_uprev.py --rust_version 1.45.0 \
+                   --state_file /tmp/state-file.json
+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.
+
+See `--help` for all available options.
+"""
+
+# pylint: disable=cros-logging-import
+
+import argparse
+import pathlib
+import json
+import logging
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, T, Tuple
+
+from llvm_tools import chroot, git
+
+
+def get_command_output(command: List[str]) -> str:
+  return subprocess.check_output(command, encoding='utf-8').strip()
+
+
+class RustVersion(NamedTuple):
+  """NamedTuple represents a Rust version"""
+  major: int
+  minor: int
+  patch: int
+
+  def __str__(self):
+    return f'{self.major}.{self.minor}.{self.patch}'
+
+  @staticmethod
+  def parse_from_ebuild(ebuild_name: str) -> 'RustVersion':
+    input_re = re.compile(r'^rust-'
+                          r'(?P<major>\d+)\.'
+                          r'(?P<minor>\d+)\.'
+                          r'(?P<patch>\d+)'
+                          r'\.ebuild$')
+    m = input_re.match(ebuild_name)
+    assert m, f'failed to parse {ebuild_name!r}'
+    return RustVersion(
+        int(m.group('major')), int(m.group('minor')), int(m.group('patch')))
+
+  @staticmethod
+  def parse(x: str) -> 'RustVersion':
+    input_re = re.compile(r'^(?:rust-)?'
+                          r'(?P<major>\d+)\.'
+                          r'(?P<minor>\d+)\.'
+                          r'(?P<patch>\d+)'
+                          r'(?:.ebuild)?$')
+    m = input_re.match(x)
+    assert m, f'failed to parse {x!r}'
+    return RustVersion(
+        int(m.group('major')), int(m.group('minor')), int(m.group('patch')))
+
+
+def parse_stage0_file(new_version: RustVersion) -> Tuple[str, str, str]:
+  # Find stage0 date, rustc and cargo
+  stage0_file = get_command_output([
+      'curl', '-f', 'https://raw.githubusercontent.com/rust-lang/rust/'
+      f'{new_version}/src/stage0.txt'
+  ])
+  regexp = re.compile(r'date:\s*(?P<date>\d+-\d+-\d+)\s+'
+                      r'rustc:\s*(?P<rustc>\d+\.\d+\.\d+)\s+'
+                      r'cargo:\s*(?P<cargo>\d+\.\d+\.\d+)')
+  m = regexp.search(stage0_file)
+  assert m, 'failed to parse stage0.txt file'
+  stage0_date, stage0_rustc, stage0_cargo = m.groups()
+  logging.info('Found stage0 file has date: %s, rustc: %s, cargo: %s',
+               stage0_date, stage0_rustc, stage0_cargo)
+  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(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)
+    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,
+                 new_version: RustVersion) -> None:
+  patch_path = os.path.join(rust_path, 'files')
+  for f in os.listdir(patch_path):
+    if f'rust-{current_version}' not in f:
+      continue
+    logging.info('Rename patch %s to new version', f)
+    new_name = f.replace(str(current_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)
+
+
+def rename_ebuild(rust_path: str, old_version: RustVersion,
+                  current_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'))
+  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')
+
+
+def update_ebuild(ebuild_file: str, stage0_info: Tuple[str, str, str]) -> None:
+  stage0_date, stage0_rustc, stage0_cargo = stage0_info
+  with open(ebuild_file, encoding='utf-8') as f:
+    contents = f.read()
+  # Update STAGE0_DATE in the ebuild
+  stage0_date_re = re.compile(r'STAGE0_DATE="(\d+-\d+-\d+)"')
+  if not stage0_date_re.search(contents):
+    raise RuntimeError('STAGE0_DATE not found in rust ebuild')
+  new_contents = stage0_date_re.sub(f'STAGE0_DATE="{stage0_date}"', contents)
+
+  # Update STAGE0_VERSION in the ebuild
+  stage0_rustc_re = re.compile(r'STAGE0_VERSION="[^"]*"')
+  if not stage0_rustc_re.search(new_contents):
+    raise RuntimeError('STAGE0_VERSION not found in rust ebuild')
+  new_contents = stage0_rustc_re.sub(f'STAGE0_VERSION="{stage0_rustc}"',
+                                     new_contents)
+
+  # Update STAGE0_VERSION_CARGO in the ebuild
+  stage0_cargo_re = re.compile(r'STAGE0_VERSION_CARGO="[^"]*"')
+  if not stage0_cargo_re.search(new_contents):
+    raise RuntimeError('STAGE0_VERSION_CARGO not found in rust ebuild')
+  new_contents = stage0_cargo_re.sub(f'STAGE0_VERSION_CARGO="{stage0_cargo}"',
+                                     new_contents)
+  with open(ebuild_file, 'w', encoding='utf-8') as f:
+    f.write(new_contents)
+  logging.info(
+      'Rust ebuild file has STAGE0_DATE, STAGE0_VERSION, STAGE0_VERSION_CARGO '
+      'updated to %s, %s, %s respectively', stage0_date, stage0_rustc,
+      stage0_cargo)
+  return ebuild_file
+
+
+def flip_mirror_in_ebuild(ebuild_file: str, add: bool) -> None:
+  restrict_re = re.compile(
+      r'(?P<before>RESTRICT=")(?P<values>"[^"]*"|.*)(?P<after>")')
+  with open(ebuild_file, encoding='utf-8') as f:
+    contents = f.read()
+  m = restrict_re.search(contents)
+  assert m, 'failed to find RESTRICT variable in Rust ebuild'
+  values = m.group('values')
+  if add:
+    if 'mirror' in values:
+      return
+    values += ' mirror'
+  else:
+    if 'mirror' not in values:
+      return
+    values = values.replace(' mirror', '')
+  new_contents = restrict_re.sub(r'\g<before>%s\g<after>' % values, contents)
+  with open(ebuild_file, 'w', encoding='utf-8') as f:
+    f.write(new_contents)
+
+
+def rust_ebuild_command(command: str, sudo: bool = False) -> None:
+  ebuild_path_inchroot = get_command_output(['equery', 'w', 'rust'])
+  cmd = ['ebuild', ebuild_path_inchroot, command]
+  if sudo:
+    cmd = ['sudo'] + cmd
+  subprocess.check_call(cmd, stderr=subprocess.STDOUT)
+
+
+def update_manifest(ebuild_file: str) -> None:
+  logging.info('Added "mirror" to RESTRICT to Rust ebuild')
+  flip_mirror_in_ebuild(ebuild_file, add=True)
+  rust_ebuild_command('manifest')
+  logging.info('Removed "mirror" to RESTRICT from Rust ebuild')
+  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:
+  package_file = os.path.join(
+      os.path.dirname(ebuild_file),
+      '../../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)
+  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,
+                        new_version: RustVersion) -> None:
+  virtual_rust_dir = os.path.join(
+      os.path.dirname(ebuild_file), '../../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)
+
+
+def upload_to_localmirror(tempdir: str, rust_version: RustVersion) -> None:
+  tarfile_name = f'rustc-{rust_version}-src.tar.gz'
+  rust_src = f'https://static.rust-lang.org/dist/{tarfile_name}'
+  logging.info('Downloading Rust from %s', rust_src)
+  gsutil_location = f'gs://chromeos-localmirror/distfiles/{tarfile_name}'
+
+  local_file = os.path.join(tempdir, tarfile_name)
+  subprocess.check_call(['curl', '-f', '-o', local_file, rust_src])
+  # Since we are using `-n` to skip an item if it already exists, there's no
+  # need to check if the file exists on GS bucket or not.
+  subprocess.check_call(
+      ['gsutil', 'cp', '-n', '-a', 'public-read', local_file, gsutil_location])
+
+
+def perform_step(state_file: pathlib.Path,
+                 tmp_state_file: pathlib.Path,
+                 completed_steps: Dict[str, Any],
+                 step_name: str,
+                 step_fn: Callable[[], T],
+                 result_from_json: Optional[Callable[[Any], T]] = None,
+                 result_to_json: Optional[Callable[[T], Any]] = None) -> T:
+  if step_name in completed_steps:
+    logging.info('Skipping previously completed step %s', step_name)
+    if result_from_json:
+      return result_from_json(completed_steps[step_name])
+    return completed_steps[step_name]
+
+  logging.info('Running step %s', step_name)
+  val = step_fn()
+  logging.info('Step %s complete', step_name)
+  if result_to_json:
+    completed_steps[step_name] = result_to_json(val)
+  else:
+    completed_steps[step_name] = val
+
+  with tmp_state_file.open('w', encoding='utf-8') as f:
+    json.dump(completed_steps, f, indent=4)
+  tmp_state_file.rename(state_file)
+  return val
+
+
+def main():
+  if not chroot.InChroot():
+    raise RuntimeError('This script must be executed inside chroot')
+
+  logging.basicConfig(level=logging.INFO)
+
+  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()
+
+  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')
+
+  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)
+
+  try:
+    with state_file.open(encoding='utf-8') as f:
+      completed_steps = json.load(f)
+  except FileNotFoundError:
+    completed_steps = {}
+
+  def run_step(
+      step_name: str,
+      step_fn: Callable[[], T],
+      result_from_json: Optional[Callable[[Any], T]] = None,
+      result_to_json: Optional[Callable[[T], Any]] = None,
+  ) -> T:
+    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 __name__ == '__main__':
+  sys.exit(main())
diff --git a/rust_tools/rust_uprev_test.py b/rust_tools/rust_uprev_test.py
new file mode 100755
index 0000000..28b23bf
--- /dev/null
+++ b/rust_tools/rust_uprev_test.py
@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Copyright 2020 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""Tests for rust_uprev.py"""
+
+# pylint: disable=cros-logging-import
+import os
+import shutil
+import subprocess
+import unittest
+from unittest import mock
+
+import rust_uprev
+from llvm_tools import git
+
+
+class RustVersionTest(unittest.TestCase):
+  """Tests for RustVersion class"""
+
+  def test_str(self):
+    obj = rust_uprev.RustVersion(major=1, minor=2, patch=3)
+    self.assertEqual(str(obj), '1.2.3')
+
+  def test_parse_version_only(self):
+    expected = rust_uprev.RustVersion(major=1, minor=2, patch=3)
+    actual = rust_uprev.RustVersion.parse('1.2.3')
+    self.assertEqual(expected, actual)
+
+  def test_parse_ebuild_name(self):
+    expected = rust_uprev.RustVersion(major=1, minor=2, patch=3)
+    actual = rust_uprev.RustVersion.parse_from_ebuild('rust-1.2.3.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')
+    self.assertEqual("failed to parse 'invalid-rust-1.2.3'",
+                     str(context.exception))
+
+
+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')
+  @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)
+    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.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_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))
+    actual = rust_uprev.prepare_uprev_from_json(json_result)
+    self.assertEqual(expected, actual)
+
+
+class UpdateEbuildTest(unittest.TestCase):
+  """Tests for update_ebuild step in rust_uprev"""
+  ebuild_file_before = """
+    STAGE0_DATE="2019-01-01"
+    STAGE0_VERSION="any.random.(number)"
+    STAGE0_VERSION_CARGO="0.0.0"
+    """
+  ebuild_file_after = """
+    STAGE0_DATE="2020-01-01"
+    STAGE0_VERSION="1.1.1"
+    STAGE0_VERSION_CARGO="0.1.0"
+    """
+
+  def test_success(self):
+    mock_open = mock.mock_open(read_data=self.ebuild_file_before)
+    ebuild_file = '/path/to/rust/rust-1.3.5.ebuild'
+    with mock.patch('builtins.open', mock_open):
+      rust_uprev.update_ebuild(ebuild_file, ('2020-01-01', '1.1.1', '0.1.0'))
+    mock_open.return_value.__enter__().write.assert_called_once_with(
+        self.ebuild_file_after)
+
+  def test_fail_when_ebuild_misses_a_variable(self):
+    ebuild_file = 'STAGE0_DATE="2019-01-01"'
+    mock_open = mock.mock_open(read_data=ebuild_file)
+    ebuild_file = '/path/to/rust/rust-1.3.5.ebuild'
+    with mock.patch('builtins.open', mock_open):
+      with self.assertRaises(RuntimeError) as context:
+        rust_uprev.update_ebuild(ebuild_file, ('2020-01-01', '1.1.1', '0.1.0'))
+    self.assertEqual('STAGE0_VERSION not found in rust ebuild',
+                     str(context.exception))
+
+
+class UpdateManifestTest(unittest.TestCase):
+  """Tests for update_manifest step in rust_uprev"""
+
+  # pylint: disable=protected-access
+  def _run_test_flip_mirror(self, before, after, add, expect_write):
+    mock_open = mock.mock_open(read_data=f'RESTRICT="{before}"')
+    with mock.patch('builtins.open', mock_open):
+      rust_uprev.flip_mirror_in_ebuild('', add=add)
+    if expect_write:
+      mock_open.return_value.__enter__().write.assert_called_once_with(
+          f'RESTRICT="{after}"')
+
+  def test_add_mirror_in_ebuild(self):
+    self._run_test_flip_mirror(
+        before='variable1 variable2',
+        after='variable1 variable2 mirror',
+        add=True,
+        expect_write=True)
+
+  def test_remove_mirror_in_ebuild(self):
+    self._run_test_flip_mirror(
+        before='variable1 variable2 mirror',
+        after='variable1 variable2',
+        add=False,
+        expect_write=True)
+
+  def test_add_mirror_when_exists(self):
+    self._run_test_flip_mirror(
+        before='variable1 variable2 mirror',
+        after='variable1 variable2 mirror',
+        add=True,
+        expect_write=False)
+
+  def test_remove_mirror_when_not_exists(self):
+    self._run_test_flip_mirror(
+        before='variable1 variable2',
+        after='variable1 variable2',
+        add=False,
+        expect_write=False)
+
+  @mock.patch.object(rust_uprev, 'flip_mirror_in_ebuild')
+  @mock.patch.object(rust_uprev, 'rust_ebuild_command')
+  def test_update_manifest(self, mock_run, mock_flip):
+    ebuild_file = '/path/to/rust/rust-1.1.1.ebuild'
+    rust_uprev.update_manifest(ebuild_file)
+    mock_run.assert_called_once_with('manifest')
+    mock_flip.assert_has_calls(
+        [mock.call(ebuild_file, add=True),
+         mock.call(ebuild_file, add=False)])
+
+
+class RustUprevOtherTests(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,
+                                    'rust-{self.new_version}.ebuild')
+
+  @mock.patch.object(rust_uprev, 'get_command_output')
+  def test_parse_stage0_file(self, mock_get):
+    stage0_file = """
+    unrelated stuff before
+    date: 2020-01-01
+    rustc: 1.1.1
+    cargo: 0.1.0
+    unrelated stuff after
+    """
+    mock_get.return_value = stage0_file
+    expected = '2020-01-01', '1.1.1', '0.1.0'
+    rust_version = rust_uprev.RustVersion(1, 2, 3)
+    actual = rust_uprev.parse_stage0_file(rust_version)
+    self.assertEqual(expected, actual)
+    mock_get.assert_called_once_with([
+        'curl', '-f', 'https://raw.githubusercontent.com/rust-lang/rust/'
+        f'{rust_version}/src/stage0.txt'
+    ])
+
+  @mock.patch.object(shutil, 'copyfile')
+  @mock.patch.object(os, 'listdir')
+  @mock.patch.object(subprocess, 'check_call')
+  def test_copy_patches(self, mock_call, mock_ls, mock_copy):
+    mock_ls.return_value = [
+        f'rust-{self.old_version}-patch-1.patch',
+        f'rust-{self.old_version}-patch-2-old.patch',
+        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)
+    mock_copy.assert_has_calls([
+        mock.call(
+            os.path.join(self.rust_path, 'files',
+                         f'rust-{self.current_version}-patch-1.patch'),
+            os.path.join(self.rust_path, 'files',
+                         f'rust-{self.new_version}-patch-1.patch'),
+        ),
+        mock.call(
+            os.path.join(self.rust_path, 'files',
+                         f'rust-{self.current_version}-patch-2-new.patch'),
+            os.path.join(self.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.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)
+    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)
+
+  @mock.patch.object(os.path, 'exists', return_value=True)
+  @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'))
+
+  @mock.patch.object(subprocess, 'check_call')
+  def test_upload_to_localmirror(self, mock_call):
+    tempdir = '/tmp/any/dir'
+    rust_uprev.upload_to_localmirror(tempdir, self.new_version)
+
+    tarfile_name = f'rustc-{self.new_version}-src.tar.gz'
+    rust_src = f'https://static.rust-lang.org/dist/{tarfile_name}'
+    gsurl = f'gs://chromeos-localmirror/distfiles/{tarfile_name}'
+    local_file = os.path.join(tempdir, tarfile_name)
+    mock_call.assert_has_calls([
+        mock.call(['curl', '-f', '-o', local_file, rust_src]),
+        mock.call(
+            ['gsutil', 'cp', '-n', '-a', 'public-read', local_file, gsurl])
+    ])
+
+
+if __name__ == '__main__':
+  unittest.main()