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

chrome/test/mini_installer/test_installer.py

Index: chrome/test/mini_installer/test_installer.py
diff --git a/chrome/test/mini_installer/test_installer.py b/chrome/test/mini_installer/test_installer.py
index c0c9a91e3b7314954d49ff6a0fb09ffef640403e..984d7a1de8e063a1f339ecd381fd40aad987cded 100644
--- a/chrome/test/mini_installer/test_installer.py
+++ b/chrome/test/mini_installer/test_installer.py
@@ -19,10 +19,12 @@ import sys
 import time
 import traceback
 import unittest
-import _winreg
 
+from state_walker import StateWalker
 from variable_expander import VariableExpander
-import verifier_runner
+
+import cleaner_visitor
+import verifier_visitor
 
 
 def LogMessage(message):
@@ -58,7 +60,7 @@ class Config:
 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:
@@ -67,6 +69,8 @@ class InstallerTest(unittest.TestCase):
           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
@@ -74,8 +78,9 @@ class InstallerTest(unittest.TestCase):
     self._config = config
     self._variable_expander = variable_expander
     self._quiet = quiet
-    self._verifier_runner = verifier_runner.VerifierRunner()
+    self._verifier = StateWalker(verifier_visitor.VerifierVisitor())
     self._clean_on_teardown = True
+    self._clean_state = clean_state
 
   def __str__(self):
     """Returns a string representing the test case.
@@ -122,7 +127,8 @@ class InstallerTest(unittest.TestCase):
   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.
@@ -143,8 +149,7 @@ class InstallerTest(unittest.TestCase):
     if not self._quiet:
       LogMessage('Verifying state %s' % state)
     try:
-      self._verifier_runner.VerifyAll(self._config.states[state],
-                                      self._variable_expander)
+      self._verifier.Walk(self._variable_expander, 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.
@@ -168,30 +173,7 @@ def RunCommand(command, variable_expander):
     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:
@@ -199,34 +181,31 @@ def RunCleanCommand(force_clean, variable_expander):
         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:
+    StateWalker(cleaner_visitor.CleanerVisitor()).Walk(
+        variable_expander, clean_state)
 
 def MergePropertyDictionaries(current_property, new_property):
   """Merges the new property dictionary into the current property dictionary.
@@ -323,6 +302,7 @@ def ParseConfigFile(filename, variable_expander):
     config.actions[action_name] = action_command
   return config
 
+DEFAULT_CLEAN_STATE = 'clean'
 
 def main():
   parser = argparse.ArgumentParser()
@@ -333,6 +313,9 @@ def main():
                       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',
@@ -340,7 +323,8 @@ def main():
   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')
@@ -361,7 +345,8 @@ def main():
                                        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__,
@@ -369,7 +354,8 @@ def main():
                               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)