[chromeos] Simplify chromeos device policy tests

chrome/test/pyautolib/policy_base.py

 # Copyright (c) 2012 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.
 
 """Base class for tests that need to update the policies enforced by Chrome.
 
 Subclasses can call SetUserPolicy (ChromeOS, Linux, Windows) and
 SetDevicePolicy (ChromeOS only) with a dictionary of the policies to install.
 
 The current implementation depends on the platform. The implementations might
@@ skipping 84 matching lines @@
     # TODO(bartfab): Extend the C++ wrapper that starts the mock DMServer so
     # that an owner can be passed in. Without this, the server will assume that
     # the owner is user@example.com and for consistency, so must we.
     owner = 'user@example.com'
     # Subclasses may override these credentials to fake enrollment into another
     # mode or use different device and machine IDs.
     mode = 'enterprise'
     device_id = string.upper(str(uuid.uuid4()))
     machine_id = 'CROSTEST'
 
+    _auth_server = None
+    _dns_server = None
+
   @staticmethod
   def _Call(command, check=False):
     """Invokes a subprocess and optionally asserts the return value is zero."""
     with open(os.devnull, 'w') as devnull:
       if check:
         return subprocess.check_call(command.split(' '), stdout=devnull)
       else:
         return subprocess.call(command.split(' '), stdout=devnull)
 
   def _WriteFile(self, path, content):
@@ skipping 199 matching lines @@
     # Restart session manager and Chrome.
     self._StartSessionManagerAndChrome()
 
   def _tearDownWithSessionManagerStopped(self):
     """Resets the test environment after stopping the session manager."""
     assert self.IsChromeOS()
     logging.debug('Stopping session manager')
     cros_ui.stop(allow_fail=True)
 
     # Stop mock GAIA server.
-    self._auth_server.stop()
+    if self._auth_server:
+      self._auth_server.stop()
 
     # Reenable TPM if present.
     if os.path.exists(TPM_SYSFS_PATH):
       self._Call('umount %s' % os.path.realpath(TPM_SYSFS_PATH))
 
     # Clear install attributes and restart cryptohomed to pick up the change.
     self._ClearInstallAttributesOnChromeOS()
 
     # Stop mock DNS server.
-    self._dns_server.stop()
+    if self._dns_server:
+      self._dns_server.stop()
 
     # Stop mock DMServer.
     self.StopHTTPServer(self._http_server)
 
     # Clear the policy served.
     pyauto_utils.RemovePath(self._temp_data_dir)
 
     # Remove the device policy blob.
     self._RemoveIfExists(constants.OWNER_KEY_FILE)
     self._RemoveIfExists(constants.SIGNED_POLICY_FILE)
 
     # Remove any existing vaults.
     self.RemoveAllCryptohomeVaultsOnChromeOS()
 
     # Restart session manager and Chrome.
     self._StartSessionManagerAndChrome()
 
   def setUp(self):
     """Sets up the platform for policy testing.
 
     On ChromeOS, part of the setup involves restarting the session manager to
     inject an initial device policy blob.
     """
     if self.IsChromeOS():
       # Perform the remainder of the setup with the device manager stopped.
-      self.WaitForSessionManagerRestart(
-          self._SetUpWithSessionManagerStopped)
+      try:
+        self.WaitForSessionManagerRestart(
+            self._SetUpWithSessionManagerStopped)
+      except:
+        # Destroy the non re-entrant services.
+        if self._auth_server:
+          self._auth_server.stop()
+        if self._dns_server:
+          self._dns_server.stop()
+        raise
 
     pyauto.PyUITest.setUp(self)
     self._branding = self.GetBrowserInfo()['properties']['branding']
 
   def tearDown(self):
     """Cleans up the policies and related files created in tests."""
     if self.IsChromeOS():
       # Perform the cleanup with the device manager stopped.
       self.WaitForSessionManagerRestart(self._tearDownWithSessionManagerStopped)
     else:
@@ skipping 189 matching lines @@
     Args:
       device_policy: The device policy to set. None clears it.
       refresh: If True, Chrome will refresh and apply the new policy.
                Requires Chrome to be alive for it.
     """
     assert self.IsChromeOS()
     self._device_policy = device_policy or {}
     self._WritePolicyOnChromeOS()
     if refresh:
       self.RefreshPolicies()