[autotest] Update login_OwnershipApi to use protobufs

client/site_tests/login_OwnershipApi/login_OwnershipApi.py

Index: client/site_tests/login_OwnershipApi/login_OwnershipApi.py
diff --git a/client/site_tests/login_OwnershipApi/login_OwnershipApi.py b/client/site_tests/login_OwnershipApi/login_OwnershipApi.py
index feda351fd0e7ded91bce7a84b78302660f856e1b..044fb9020f0b8daaf6c91ded17d5cb2605100d41 100644
--- a/client/site_tests/login_OwnershipApi/login_OwnershipApi.py
+++ b/client/site_tests/login_OwnershipApi/login_OwnershipApi.py
@@ -7,9 +7,10 @@ import dbus.glib
 import gobject
 import logging
 import os
+import sys
 import tempfile
 
-from autotest_lib.client.bin import test
+from autotest_lib.client.bin import test, utils
 from autotest_lib.client.common_lib import autotemp, error
 from autotest_lib.client.cros import constants, cros_ui, cryptohome, login
 from autotest_lib.client.cros import ownership
@@ -20,16 +21,25 @@ class login_OwnershipApi(test.test):
 
     _testuser = 'cryptohometest@chromium.org'
     _testpass = 'testme'
-    _testpolicydata = 'hooberbloob'
+    _poldata = 'hooberbloob'
 
     _tempdir = None
 
-    def initialize(self):
+    def setup(self):
+        os.chdir(self.srcdir)
+        utils.make('OUT_DIR=.')
+
+
+    def __unlink(self, filename):
         try:
-            os.unlink(constants.OWNER_KEY_FILE)
-            os.unlink(constants.SIGNED_PREFERENCES_FILE)
+            os.unlink(filename)
         except (IOError, OSError) as error:
             logging.info(error)
+
+    def initialize(self):
+        self.__unlink(constants.OWNER_KEY_FILE)
+        self.__unlink(constants.SIGNED_PREFERENCES_FILE)
+        self.__unlink(constants.SIGNED_POLICY_FILE)
         login.refresh_login_screen()
         cryptohome.remove_vault(self._testuser)
         cryptohome.mount_vault(self._testuser, self._testpass, create=True)
@@ -63,6 +73,9 @@ class login_OwnershipApi(test.test):
     def run_once(self):
         keyfile = ownership.generate_and_register_owner_keypair(self._testuser,
                                                                 self._testpass)
+        # Pull in protobuf definitions.
+        sys.path.append(self.srcdir)
+        from device_management_backend_pb2 import PolicyFetchResponse
 
         # open DBus connection to session_manager
         bus = dbus.SystemBus()
@@ -70,18 +83,21 @@ class login_OwnershipApi(test.test):
                                '/org/chromium/SessionManager')
         sm = dbus.Interface(proxy, 'org.chromium.SessionManagerInterface')
 
-        sig = ownership.sign(keyfile, self._testuser)
-        sm.Whitelist(self._testuser, dbus.ByteArray(sig))
-        wl_sig = sm.CheckWhitelist(self._testuser, byte_arrays=True)
-        if sig != wl_sig:
-            raise error.TestFail("CheckWhitelist signature mismatch")
-
-        sm.Unwhitelist(self._testuser, dbus.ByteArray(sig))
-        try:
-            sm.CheckWhitelist(self._testuser)
-            raise error.TestFail("Should not have found user in whitelist!")
-        except dbus.DBusException as e:
-            logging.debug(e)
+        policy_proto = PolicyFetchResponse()
+        policy_proto.policy_data = self._poldata
+        policy_proto.policy_data_signature = ownership.sign(keyfile,
+                                                            self._poldata)
+        sm.StorePolicy(dbus.ByteArray(policy_proto.SerializeToString()),
+                       byte_arrays=True,
+                       reply_handler=self.__log_and_stop,
+                       error_handler=self.__log_err_and_stop)
+
+        self._loop = gobject.MainLoop()
+        self._loop.run()
+
+        retrieved_policy = sm.RetrievePolicy(byte_arrays=True)
+        if retrieved_policy != policy_proto.SerializeToString():
+            raise error.TestFail('Policy should not be %s' % retrieved_policy)
 
 
     def cleanup(self):