Listen for ownership-taken signal on DBus

Many tests that touch on device Ownership were using a
function that simply watched for the presence of policy
and owner key files to indicate that the session_manager
had completed the ownership-taking process on the sign-in
of the first user. This could be tricked, though, in the
case of the UI being restarted between the test clearing
these files and then checking for them again. Stop using
this and switch to watching for signals announcing that
ownership has been taken, which the session_manager has
sent all along anyhow.

BUG=chromium:355664
TEST=login_*Ownership*

Change-Id: I54000c426cfb1aef8fd1518f9cdadec2b42c6500
Reviewed-on: https://chromium-review.googlesource.com/194142
Reviewed-by: Chris Masone <cmasone@chromium.org>
Commit-Queue: Chris Masone <cmasone@chromium.org>
Tested-by: Chris Masone <cmasone@chromium.org>
diff --git a/client/common_lib/cros/policy.py b/client/common_lib/cros/policy.py
index 09a4e23..b039a7a 100644
--- a/client/common_lib/cros/policy.py
+++ b/client/common_lib/cros/policy.py
@@ -2,11 +2,12 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import dbus, sys
+import dbus, gobject, sys
 
 import common
 from autotest_lib.client.common_lib import error
-from autotest_lib.client.cros import login, ownership
+from autotest_lib.client.common_lib.cros import session_manager
+from autotest_lib.client.cros import ownership
 
 
 """Utility class for tests that generate, push and fetch policies.
@@ -165,8 +166,10 @@
 
     @raises error.TestFail if policy push failed.
     """
+    listener = session_manager.OwnershipSignalListener(gobject.MainLoop())
+    listener.listen_for_new_policy()
     sm.StorePolicy(dbus.ByteArray(policy_string), byte_arrays=True)
-    login.wait_for_ownership()
+    listener.wait_for_signals(desc='Policy push.')
 
     retrieved_policy = sm.RetrievePolicy(byte_arrays=True)
     if retrieved_policy != policy_string:
diff --git a/client/common_lib/cros/session_manager.py b/client/common_lib/cros/session_manager.py
index 5457b02..4b8d6a5 100644
--- a/client/common_lib/cros/session_manager.py
+++ b/client/common_lib/cros/session_manager.py
@@ -114,20 +114,36 @@
         @param g_mail_loop: glib main loop object.
         """
         super(OwnershipSignalListener, self).__init__(g_main_loop)
+        self._listen_for_new_key = False
         self._got_new_key = False
+        self._listen_for_new_policy = False
         self._got_new_policy = False
 
+
     def listen_for_new_key_and_policy(self):
         """Set to listen for signals indicating new owner key and device policy.
         """
+        self._listen_for_new_key = self._listen_for_new_policy = True
         self.listen_to_signal(self.__handle_new_key, 'SetOwnerKeyComplete')
         self.listen_to_signal(self.__handle_new_policy,
                               'PropertyChangeComplete')
+        self.reset_signal_state()
+
+
+    def listen_for_new_policy(self):
+        """Set to listen for signal indicating new device policy.
+        """
+        self._listen_for_new_key = False
+        self._listen_for_new_policy = True
+        self.listen_to_signal(self.__handle_new_policy,
+                              'PropertyChangeComplete')
+        self.reset_signal_state()
 
 
     def reset_signal_state(self):
         """Resets internal signal tracking state."""
-        self._got_new_policy = self._got_new_key = False
+        self._got_new_key = not self._listen_for_new_key
+        self._got_new_policy = not self._listen_for_new_policy
 
 
     def all_signals_received(self):
diff --git a/client/cros/login.py b/client/cros/login.py
index 9ec8805..9637c5d 100644
--- a/client/cros/login.py
+++ b/client/cros/login.py
@@ -25,12 +25,31 @@
     """Checks the log watched by |log_reader| to see if a crash was reported
     for |process|.
 
-    Returns True if so, False if not.
+    @param process: process name to look for.
+    @param log_reader: LogReader object set up to watch appropriate log file.
+
+    @return: True if so, False if not.
     """
     return log_reader.can_find('Received crash notification for %s' % process)
 
 
 def wait_for_condition(condition, timeout_msg, timeout, process, crash_msg):
+    """Wait for callable |condition| to return true, while checking for crashes.
+
+    Poll for |condition| to become true, for |timeout| seconds. If the timeout
+    is reached, check to see if |process| crashed while we were polling.
+    If so, raise CrashError(crash_msg). If not, raise TimeoutError(timeout_msg).
+
+    @param condition: a callable to poll on.
+    @param timeout_msg: message to put in TimeoutError before raising.
+    @param timeout: float number of seconds to poll on |condition|.
+    @param process: process name to watch for crashes while polling.
+    @param crash_msg: message to put in CrashError if polling failed and
+                      |process| crashed.
+
+    @raise: TimeoutError if timeout is reached.
+    @raise: CrashError if process crashed and the condition never fired.
+    """
     # Mark /var/log/messages now; we'll run through all subsequent log
     # messages if we couldn't start chrome to see if the browser crashed.
     log_reader = cros_logging.LogReader()
@@ -52,11 +71,9 @@
 def wait_for_browser(timeout=cros_ui.RESTART_UI_TIMEOUT):
     """Wait until a Chrome process is running.
 
-    Args:
-        timeout: float number of seconds to wait
+    @param timeout: float number of seconds to wait.
 
-    Raises:
-        TimeoutError: Chrome didn't start before timeout
+    @raise: TimeoutError: Chrome didn't start before timeout.
     """
     wait_for_condition(
         lambda: os.system('pgrep ^%s$ >/dev/null' % constants.BROWSER) == 0,
@@ -69,15 +86,12 @@
 def wait_for_browser_exit(crash_msg, timeout=cros_ui.RESTART_UI_TIMEOUT):
     """Wait for the Chrome process to exit.
 
-    Args:
-        crash_msg: Error message to include if Chrome crashed.
-        timeout: float number of seconds to wait
+    @param crash_msg: Error message to include if Chrome crashed.
+    @param timeout: float number of seconds to wait.
 
-    Returns:
-        True if Chrome exited; False otherwise.
+    @return: True if Chrome exited; False otherwise.
 
-    Raises:
-        CrashError: Chrome crashed while we were waiting.
+    @raise: CrashError: Chrome crashed while we were waiting.
     """
     try:
       wait_for_condition(
@@ -94,11 +108,10 @@
 def wait_for_cryptohome(user, timeout=cros_ui.RESTART_UI_TIMEOUT):
     """Wait until cryptohome is mounted.
 
-    Args:
-        timeout: float number of seconds to wait
+    @param user: the user whose cryptohome the caller wants to wait for.
+    @param timeout: float number of seconds to wait.
 
-    Raises:
-        TimeoutError: cryptohome wasn't mounted before timeout
+    @raise: TimeoutError: cryptohome wasn't mounted before timeout
     """
     wait_for_condition(
         condition=lambda: cryptohome.is_vault_mounted(user),
@@ -109,6 +122,14 @@
 
 
 def wait_for_ownership(timeout=constants.DEFAULT_OWNERSHIP_TIMEOUT):
+    """Wait until device owner key file exists on disk.
+
+    @param timeout: float number of seconds to wait.
+
+    @raise: TimeoutError: file didn't appear before timeout.
+    """
+    if os.access(constants.OWNER_KEY_FILE, os.F_OK):
+        raise error.TestError('Device is already owned!')
     wait_for_condition(
         condition=lambda: os.access(constants.OWNER_KEY_FILE, os.F_OK),
         timeout_msg='Timed out waiting for ownership',
diff --git a/client/site_tests/login_OwnershipNotRetaken/login_OwnershipNotRetaken.py b/client/site_tests/login_OwnershipNotRetaken/login_OwnershipNotRetaken.py
index 96485f1..e417c1a 100644
--- a/client/site_tests/login_OwnershipNotRetaken/login_OwnershipNotRetaken.py
+++ b/client/site_tests/login_OwnershipNotRetaken/login_OwnershipNotRetaken.py
@@ -2,13 +2,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import hashlib, os
+import gobject, hashlib, os
 from dbus.mainloop.glib import DBusGMainLoop
 
 from autotest_lib.client.bin import test
 from autotest_lib.client.common_lib import error
-from autotest_lib.client.common_lib.cros import chrome
-from autotest_lib.client.cros import constants, cryptohome, login
+from autotest_lib.client.common_lib.cros import chrome, session_manager
+from autotest_lib.client.cros import constants, cryptohome, ownership
 
 
 class login_OwnershipNotRetaken(test.test):
@@ -19,10 +19,21 @@
     _TEST_PASS = 'testme'
 
 
+    def initialize(self):
+        super(login_OwnershipNotRetaken, self).initialize()
+        # Start clean, wrt ownership and the desired user.
+        ownership.restart_ui_to_clear_ownership_files()
+
+        bus_loop = DBusGMainLoop(set_as_default=True)
+        self._cryptohome_proxy = cryptohome.CryptohomeProxy(bus_loop)
+
+
     def run_once(self):
+        listener = session_manager.OwnershipSignalListener(gobject.MainLoop())
+        listener.listen_for_new_key_and_policy()
         # Sign in. Sign out happens automatically when cr goes out of scope.
         with chrome.Chrome() as cr:
-            login.wait_for_ownership()
+            listener.wait_for_signals(desc='Owner settings written to disk.')
 
         key = open(constants.OWNER_KEY_FILE, 'rb')
         hash = hashlib.md5(key.read())
@@ -48,5 +59,4 @@
 
     def cleanup(self):
         super(login_OwnershipNotRetaken, self).cleanup()
-        bus_loop = DBusGMainLoop(set_as_default=True)
-        cryptohome.CryptohomeProxy(bus_loop).remove(self._TEST_USER)
+        self._cryptohome_proxy.remove(self._TEST_USER)
diff --git a/client/site_tests/login_OwnershipRetaken/login_OwnershipRetaken.py b/client/site_tests/login_OwnershipRetaken/login_OwnershipRetaken.py
index 94d41ab..955a130 100644
--- a/client/site_tests/login_OwnershipRetaken/login_OwnershipRetaken.py
+++ b/client/site_tests/login_OwnershipRetaken/login_OwnershipRetaken.py
@@ -34,9 +34,6 @@
         self._cryptohome_proxy.remove(ownership.TESTUSER)
 
         self._sm = session_manager.connect(bus_loop)
-        self._listener = session_manager.OwnershipSignalListener(
-                gobject.MainLoop())
-        self._listener.listen_for_new_key_and_policy()
 
 
     def run_once(self):
@@ -58,20 +55,20 @@
                                                poldata)
         policy.push_policy_and_verify(policy_string, self._sm)
 
-        self._listener.wait_for_signals(desc='Initial policy push complete.')
-
         # grab key, ensure that it's the same as the known key.
         if (utils.read_file(constants.OWNER_KEY_FILE) != pubkey):
             raise error.TestFail('Owner key should not have changed!')
 
         # Start a new session, which will trigger the re-taking of ownership.
+        listener = session_manager.OwnershipSignalListener(gobject.MainLoop())
+        listener.listen_for_new_key_and_policy()
         self._cryptohome_proxy.mount(ownership.TESTUSER,
                                      ownership.TESTPASS,
                                      create=True)
         if not self._sm.StartSession(ownership.TESTUSER, ''):
-            raise error.TestFail('Could not start session for owner')
+            raise error.TestError('Could not start session for owner')
 
-        self._listener.wait_for_signals(desc='Re-taking of ownership complete.')
+        listener.wait_for_signals(desc='Re-taking of ownership complete.')
 
         # grab key, ensure that it's different than known key
         if (utils.read_file(constants.OWNER_KEY_FILE) == pubkey):
@@ -80,7 +77,7 @@
         # RetrievePolicy, check sig against new key, check properties
         retrieved_policy = self._sm.RetrievePolicy(byte_arrays=True)
         if retrieved_policy is None:
-            raise error.TestFail('Policy not found')
+            raise error.TestError('Policy not found')
         policy.compare_policy_response(self.srcdir,
                                        retrieved_policy,
                                        owner=ownership.TESTUSER,
diff --git a/client/site_tests/login_OwnershipTaken/login_OwnershipTaken.py b/client/site_tests/login_OwnershipTaken/login_OwnershipTaken.py
index bf6e391..a7eff12 100644
--- a/client/site_tests/login_OwnershipTaken/login_OwnershipTaken.py
+++ b/client/site_tests/login_OwnershipTaken/login_OwnershipTaken.py
@@ -2,13 +2,13 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-import os, sys
+import gobject, os, sys
 from dbus.mainloop.glib import DBusGMainLoop
 
 from autotest_lib.client.bin import test, utils
 from autotest_lib.client.common_lib import error
 from autotest_lib.client.common_lib.cros import chrome, session_manager
-from autotest_lib.client.cros import constants, login, ownership
+from autotest_lib.client.cros import constants, ownership
 
 
 class login_OwnershipTaken(test.test):
@@ -22,6 +22,7 @@
 
 
     def initialize(self):
+        super(login_OwnershipTaken, self).initialize()
         ownership.restart_ui_to_clear_ownership_files()
         if (os.access(constants.OWNER_KEY_FILE, os.F_OK) or
             os.access(constants.SIGNED_POLICY_FILE, os.F_OK)):
@@ -53,8 +54,10 @@
 
     def run_once(self):
         bus_loop = DBusGMainLoop(set_as_default=True)
+        listener = session_manager.OwnershipSignalListener(gobject.MainLoop())
+        listener.listen_for_new_key_and_policy()
         with chrome.Chrome() as cr:
-            login.wait_for_ownership()
+            listener.wait_for_signals(desc='Owner settings written to disk.')
 
             sm = session_manager.connect(bus_loop)
             retrieved_policy = sm.RetrievePolicy(byte_arrays=True)