[Autotest] Adding a wait option when refreshing policies

Added a wait option when refreshing policies. When enabled, will wait up
to 1 second for the polices to change after the refresh is called. The
wait will end either when the policies are changed, or the timeout(1
second) has expired. If the timeout expires, the test will simply
continue

BUG=None
TEST=policy_ExtensionPolicy (the failing test)

Change-Id: I2599b953c24beaac7631b023403db7ccd28a58a9
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/autotest/+/1960429
Reviewed-by: Ruben Zakarian <rzakarian@chromium.org>
Tested-by: Derek Beckett <dbeckett@chromium.org>
Commit-Queue: Derek Beckett <dbeckett@chromium.org>
diff --git a/client/cros/enterprise/enterprise_policy_base.py b/client/cros/enterprise/enterprise_policy_base.py
index 6626d8f..eb03bbf 100755
--- a/client/cros/enterprise/enterprise_policy_base.py
+++ b/client/cros/enterprise/enterprise_policy_base.py
@@ -320,7 +320,7 @@
 
         """
         self.pol_manager.remove_policy(name, policy_type, extID)
-        self.reload_policies()
+        self.reload_policies(verify_policies)
         if verify_policies:
             self.pol_manager.verify_policies()
 
@@ -423,7 +423,7 @@
                                             device=device,
                                             extension=extension,
                                             new=new)
-        self.reload_policies()
+        self.reload_policies(True)
 
     def update_policies(self, user_policies={}, suggested_user_policies={},
                         device_policies={}, extension_policies={}):
@@ -441,9 +441,16 @@
         self.add_policies(user_policies, suggested_user_policies,
                           device_policies, extension_policies, True)
 
-    def reload_policies(self):
-        """Force a policy fetch."""
-        enterprise_policy_utils.refresh_policies(self.cr.autotest_ext)
+    def reload_policies(self, wait_for_new=False):
+        """
+        Force a policy fetch.
+
+        @param wait_for_new: bool, wait up to 1 second for the policy values
+            from the API to update
+
+        """
+        enterprise_policy_utils.refresh_policies(self.cr.autotest_ext,
+                                                 wait_for_new)
 
     def verify_extension_stats(self, extension_policies, sensitive_fields=[]):
         """
diff --git a/client/cros/enterprise/enterprise_policy_utils.py b/client/cros/enterprise/enterprise_policy_utils.py
index facbcb0..a55ec2e 100644
--- a/client/cros/enterprise/enterprise_policy_utils.py
+++ b/client/cros/enterprise/enterprise_policy_utils.py
@@ -12,9 +12,12 @@
 
 """
 import json
+import time
+
 from autotest_lib.client.common_lib import error
 # Default settings for managed user policies
 
+
 def get_all_policies(autotest_ext):
     """Returns a dict of all the policies on the device."""
 
@@ -41,8 +44,25 @@
     return policy_data
 
 
-def refresh_policies(autotest_ext):
-    """Force a policy fetch."""
+def _wait_for_new_pols_after_refresh(autotest_ext):
+    """
+    Wait up to 1 second for the polices to update when refreshing.
+
+    Note: This is non-breaking, thus even if the policies do not update, once
+    the timeout expires, the function will exit.
+
+    """
+    prior = _get_pol_from_api(autotest_ext)
+    _call_refresh_policies_from_api(autotest_ext)
+    t1 = time.time()
+    while time.time() - t1 < 1:
+        curr = _get_pol_from_api(autotest_ext)
+        if curr != prior:
+            break
+
+
+def _call_refresh_policies_from_api(autotest_ext):
+    """Call the refreshEnterprisePolicies from autotestPrivate."""
     new_promise = '''new Promise(function(resolve, reject) {
         chrome.autotestPrivate.refreshEnterprisePolicies(function() {
           if (chrome.runtime.lastError) {
@@ -55,6 +75,15 @@
     autotest_ext.EvaluateJavaScript(new_promise, promise=True)
 
 
+def refresh_policies(autotest_ext, wait_for_new=False):
+    """Force a policy fetch."""
+
+    if wait_for_new:
+        _wait_for_new_pols_after_refresh(autotest_ext)
+    else:
+        _call_refresh_policies_from_api(autotest_ext)
+
+
 def _reformat_policies(policy_dict):
     """
     Reformat visually formatted dicts to type dict (and not unicode).