Cleanup machine based on the state in configuration file for mini installer test.

chrome/test/mini_installer/test_installer.py

 # Copyright 2013 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """This script tests the installer with test cases specified in the config file.
 
 For each test case, it checks that the machine states after the execution of
 each command match the expected machine states. For more details, take a look at
 the design documentation at http://goo.gl/Q0rGM6
 """
 
 import argparse
 import datetime
 import inspect
 import json
 import os
 import subprocess
 import sys
 import time
 import traceback
 import unittest
-import _winreg
 
 from variable_expander import VariableExpander
-import verifier_runner
+import visitor_runner
 
 
 def LogMessage(message):
   """Logs a message to stderr.
 
   Args:
     message: The message string to be logged.
   """
   now = datetime.datetime.now()
   frameinfo = inspect.getframeinfo(inspect.currentframe().f_back)
@@ skipping 15 matching lines @@
   """
   def __init__(self):
     self.states = {}
     self.actions = {}
     self.tests = []
 
 
 class InstallerTest(unittest.TestCase):
   """Tests a test case in the config file."""
 
-  def __init__(self, name, test, config, variable_expander, quiet):
+  def __init__(self, name, test, config, variable_expander, quiet, clean_state):
     """Constructor.
 
     Args:
       name: The name of this test.
       test: An array of alternating state names and action names, starting and
           ending with state names.
       config: The Config object.
       variable_expander: A VariableExpander object.
+      quiet: A boolean to control the test output.
+      clean_state: A string to represent the cleanup state.
     """
     super(InstallerTest, self).__init__()
     self._name = name
     self._test = test
     self._config = config
     self._variable_expander = variable_expander
     self._quiet = quiet
-    self._verifier_runner = verifier_runner.VerifierRunner()
+    self._verifier_runner = visitor_runner.VerifierRunner(variable_expander)
     self._clean_on_teardown = True
+    self._clean_state = clean_state
 
   def __str__(self):
     """Returns a string representing the test case.
 
     Returns:
       A string created by joining state names and action names together with
       ' -> ', for example, 'Test: clean -> install chrome -> chrome_installed'.
     """
     return '%s: %s\n' % (self._name, ' -> '.join(self._test))
 
@@ skipping 26 matching lines @@
       state = self._test[i + 1]
       self._VerifyState(state)
 
     # If the test makes it here, it means it was successful, because RunCommand
     # and _VerifyState throw an exception on failure.
     self._clean_on_teardown = False
 
   def tearDown(self):
     """Cleans up the machine if the test case fails."""
     if self._clean_on_teardown:
-      RunCleanCommand(True, self._variable_expander)
+      RunCleanCommand(True, self._config.states[self._clean_state],
+                      self._variable_expander)
 
   def shortDescription(self):
     """Overridden from unittest.TestCase.
 
     We return None as the short description to suppress its printing.
     The default implementation of this method returns the docstring of the
     runTest method, which is not useful since it's the same for every test case.
     The description from the __str__ method is informative enough.
     """
     return None
 
   def _VerifyState(self, state):
     """Verifies that the current machine state matches a given state.
 
     Args:
       state: A state name.
     """
     if not self._quiet:
       LogMessage('Verifying state %s' % state)
     try:
-      self._verifier_runner.VerifyAll(self._config.states[state],
-                                      self._variable_expander)
+      self._verifier_runner.Run(self._config.states[state])
     except AssertionError as e:
       # If an AssertionError occurs, we intercept it and add the state name
       # to the error message so that we know where the test fails.
       raise AssertionError("In state '%s', %s" % (state, e))
 
 
 def RunCommand(command, variable_expander):
   """Runs the given command from the current file's directory.
 
   This function throws an Exception if the command returns with non-zero exit
   status.
 
   Args:
     command: A command to run. It is expanded using Expand.
     variable_expander: A VariableExpander object.
   """
   expanded_command = variable_expander.Expand(command)
   script_dir = os.path.dirname(os.path.abspath(__file__))
   exit_status = subprocess.call(expanded_command, shell=True, cwd=script_dir)
   if exit_status != 0:
     raise Exception('Command %s returned non-zero exit status %s' % (
         expanded_command, exit_status))
 
-
-def DeleteGoogleUpdateRegistration(system_level, registry_subkey,
-                                   variable_expander):
-  """Deletes Chrome's registration with Google Update.
-
-  Args:
-    system_level: True if system-level Chrome is to be deleted.
-    registry_subkey: The pre-expansion registry subkey for the product.
-    variable_expander: A VariableExpander object.
-  """
-  root = (_winreg.HKEY_LOCAL_MACHINE if system_level
-          else _winreg.HKEY_CURRENT_USER)
-  key_name = variable_expander.Expand(registry_subkey)
-  try:
-    key_handle = _winreg.OpenKey(root, key_name, 0,
-                                 _winreg.KEY_SET_VALUE |
-                                 _winreg.KEY_WOW64_32KEY)
-    _winreg.DeleteValue(key_handle, 'pv')
-  except WindowsError:
-    # The key isn't present, so there is no value to delete.
-    pass
-
-
-def RunCleanCommand(force_clean, variable_expander):
+def RunCleanCommand(force_clean, clean_state, variable_expander):
   """Puts the machine in the clean state (i.e. Chrome not installed).
 
   Args:
     force_clean: A boolean indicating whether to force cleaning existing
         installations.
     variable_expander: A VariableExpander object.
   """
-  # A list of (system_level, product_name, product_switch, registry_subkey)
+  # A list of (product_name, product_switch)
   # tuples for the possible installed products.
   data = [
-    (False, '$CHROME_LONG_NAME', '',
-     '$CHROME_UPDATE_REGISTRY_SUBKEY'),
-    (True, '$CHROME_LONG_NAME', '--system-level',
-     '$CHROME_UPDATE_REGISTRY_SUBKEY'),
+    ('$CHROME_LONG_NAME', ''),
+    ('$CHROME_LONG_NAME', '--system-level'),
   ]
   if variable_expander.Expand('$SUPPORTS_SXS') == 'True':
-    data.append((False, '$CHROME_LONG_NAME_SXS', '',
-                 '$CHROME_UPDATE_REGISTRY_SUBKEY_SXS'))
+    data.append(('$CHROME_LONG_NAME_SXS', ''))
 
   interactive_option = '--interactive' if not force_clean else ''
-  for system_level, product_name, product_switch, registry_subkey in data:
+  for product_name, product_switch in data:
     command = ('python uninstall_chrome.py '
                '--chrome-long-name="%s" '
                '--no-error-if-absent %s %s' %
                (product_name, product_switch, interactive_option))
     try:
       RunCommand(command, variable_expander)
-    except:
+    except (Exception, OSError, ValueError):
       message = traceback.format_exception(*sys.exc_info())
       message.insert(0, 'Error cleaning up an old install with:\n')
       LogMessage(''.join(message))
-    if force_clean:
-      DeleteGoogleUpdateRegistration(system_level, registry_subkey,
-                                     variable_expander)
+
+  if force_clean:
+    visitor_runner.CleanupRunner.Create(variable_expander).Run(clean_state)
 
 
 def MergePropertyDictionaries(current_property, new_property):
   """Merges the new property dictionary into the current property dictionary.
 
   This is different from general dictionary merging in that, in case there are
   keys with the same name, we merge values together in the first level, and we
   override earlier values in the second level. For more details, take a look at
   http://goo.gl/uE0RoR
 
@@ skipping 77 matching lines @@
                                                         variable_expander),
                         config.tests)
   for state_name, state_property_filenames in config_data['states']:
     config.states[state_name] = ParsePropertyFiles(directory,
                                                    state_property_filenames,
                                                    variable_expander)
   for action_name, action_command in config_data['actions']:
     config.actions[action_name] = action_command
   return config
 
+DEFAULT_CLEAN_STATE = 'clean'
 
 def main():
   parser = argparse.ArgumentParser()
   parser.add_argument('--build-dir', default='out',
                       help='Path to main build directory (the parent of the '
                       'Release or Debug directory)')
   parser.add_argument('--target', default='Release',
                       help='Build target (Release or Debug)')
   parser.add_argument('--force-clean', action='store_true', default=False,
                       help='Force cleaning existing installations')
+  parser.add_argument('--clean-state', default='clean',
+                      help='The state that is used to cleanup the machine after'
+                            'each test case.')
   parser.add_argument('-q', '--quiet', action='store_true', default=False,
                       help='Reduce test runner output')
   parser.add_argument('--write-full-results-to', metavar='FILENAME',
                       help='Path to write the list of full results to.')
   parser.add_argument('--config', metavar='FILENAME',
                       help='Path to test configuration file')
   parser.add_argument('test', nargs='*',
-                      help='Name(s) of tests to run.')
+                      help=('Name(s) of tests to run. For example, '
+                           '__main__.InstallerTest.ChromeUserLevel'))
   args = parser.parse_args()
   if not args.config:
     parser.error('missing mandatory --config FILENAME argument')
 
   mini_installer_path = os.path.join(args.build_dir, args.target,
                                      'mini_installer.exe')
   assert os.path.exists(mini_installer_path), ('Could not find file %s' %
                                                mini_installer_path)
 
   next_version_mini_installer_path = os.path.join(
       args.build_dir, args.target, 'next_version_mini_installer.exe')
   assert os.path.exists(next_version_mini_installer_path), (
       'Could not find file %s' % next_version_mini_installer_path)
 
   suite = unittest.TestSuite()
 
   variable_expander = VariableExpander(mini_installer_path,
                                        next_version_mini_installer_path)
   config = ParseConfigFile(args.config, variable_expander)
 
-  RunCleanCommand(args.force_clean, variable_expander)
+  RunCleanCommand(args.force_clean, config.states[args.clean_state],
+                  variable_expander)
   for test in config.tests:
     # If tests were specified via |tests|, their names are formatted like so:
     test_name = '%s.%s.%s' % (InstallerTest.__module__,
                               InstallerTest.__name__,
                               test['name'])
     if not args.test or test_name in args.test:
       suite.addTest(InstallerTest(test['name'], test['traversal'], config,
-                                  variable_expander, args.quiet))
+                                  variable_expander, args.quiet,
+                                  args.clean_state))
 
   verbosity = 2 if not args.quiet else 1
   result = unittest.TextTestRunner(verbosity=verbosity).run(suite)
   if args.write_full_results_to:
     with open(args.write_full_results_to, 'w') as fp:
       json.dump(_FullResults(suite, result, {}), fp, indent=2)
       fp.write('\n')
   return 0 if result.wasSuccessful() else 1
 
 
@@ skipping 59 matching lines @@
     trie[path] = value
     return
   directory, rest = path.split(TEST_SEPARATOR, 1)
   if directory not in trie:
     trie[directory] = {}
   _AddPathToTrie(trie[directory], rest, value)
 
 
 if __name__ == '__main__':
   sys.exit(main())