Fix flakiness in ShutdownPolicyLoginTest

The test updates shutdown policy, and waits for the login shelf UI to
get updated from the state before the policy was changed. It then
verifies that the shutdown/restart button visibility has been updated.

This is racy if login shelf state can change for any reason different
than shutdown policy value - which started happening since the test was
updated to properly set-up enrollment, OOBE complete state. The login
UI lands on login screen with GAIA dialog, instead of welcome screen,
and GAIA UI is trying to update guest login button visibility.

This fixes the test to wait for the shutdown button visibility to
change. It observes login shelf UI updates until it detects that the
shutdown button has the desired visibility state.

Does the smae for lock test, even though that one does not seem flaky
for now.

Change-Id: I4b61c17c2ef750b9850bbbe68ea406b15197ec81
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1608525Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Commit-Queue: Toni Baržić <tbarzic@chromium.org>
Cr-Commit-Position: refs/heads/master@{#659074}

chrome/browser/chromeos/shutdown_policy_browsertest.cc

@@ -191,6 +191,15 @@ class ShutdownPolicyLockerTest : public ShutdownPolicyBaseTest {
     ShutdownPolicyBaseTest::TearDownOnMainThread();
   }
 
+  void WaitForShutdownButtonVisibility(bool visible) {
+    ScreenLockerTester tester;
+    int ui_update_count = tester.GetUiUpdateCount();
+    while (tester.IsLockShutdownButtonShown() != visible) {
+      tester.WaitForUiUpdate(ui_update_count);
+      ui_update_count = tester.GetUiUpdateCount();
+    }
+  }
+
  private:
   std::unique_ptr<ui::ScopedAnimationDurationScaleMode> zero_duration_mode_;
 
@@ -205,19 +214,15 @@ IN_PROC_BROWSER_TEST_F(ShutdownPolicyLockerTest, TestBasic) {
 
 IN_PROC_BROWSER_TEST_F(ShutdownPolicyLockerTest, PolicyChange) {
   ScreenLockerTester tester;
-  int ui_update_count = tester.GetUiUpdateCount();
   UpdateRebootOnShutdownPolicy(true);
   RefreshDevicePolicy();
-  tester.WaitForUiUpdate(ui_update_count);
+  WaitForShutdownButtonVisibility(false);
   EXPECT_TRUE(tester.IsLockRestartButtonShown());
-  EXPECT_FALSE(tester.IsLockShutdownButtonShown());
 
-  ui_update_count = tester.GetUiUpdateCount();
   UpdateRebootOnShutdownPolicy(false);
   RefreshDevicePolicy();
-  tester.WaitForUiUpdate(ui_update_count);
+  WaitForShutdownButtonVisibility(true);
   EXPECT_FALSE(tester.IsLockRestartButtonShown());
-  EXPECT_TRUE(tester.IsLockShutdownButtonShown());
 }
 
 class ShutdownPolicyLoginTest : public ShutdownPolicyBaseTest {
@@ -254,6 +259,15 @@ class ShutdownPolicyLoginTest : public ShutdownPolicyBaseTest {
     }
   }
 
+  void WaitForShutdownButtonVisibility(bool visible) {
+    test::LoginScreenTester tester;
+    int ui_update_count = tester.GetUiUpdateCount();
+    while (tester.IsShutdownButtonShown() != visible) {
+      tester.WaitForUiUpdate(ui_update_count);
+      ui_update_count = tester.GetUiUpdateCount();
+    }
+  }
+
  private:
   DISALLOW_COPY_AND_ASSIGN(ShutdownPolicyLoginTest);
 };
@@ -266,19 +280,15 @@ IN_PROC_BROWSER_TEST_F(ShutdownPolicyLoginTest, PolicyNotSet) {
 
 IN_PROC_BROWSER_TEST_F(ShutdownPolicyLoginTest, PolicyChange) {
   test::LoginScreenTester tester;
-  int ui_update_count = tester.GetUiUpdateCount();
   UpdateRebootOnShutdownPolicy(true);
   RefreshDevicePolicy();
-  tester.WaitForUiUpdate(ui_update_count);
+  WaitForShutdownButtonVisibility(false);
   EXPECT_TRUE(tester.IsRestartButtonShown());
-  EXPECT_FALSE(tester.IsShutdownButtonShown());
 
-  ui_update_count = tester.GetUiUpdateCount();
   UpdateRebootOnShutdownPolicy(false);
   RefreshDevicePolicy();
-  tester.WaitForUiUpdate(ui_update_count);
+  WaitForShutdownButtonVisibility(true);
   EXPECT_FALSE(tester.IsRestartButtonShown());
-  EXPECT_TRUE(tester.IsShutdownButtonShown());
 }
 
 }  // namespace chromeos