Correct UAF errors from KioskNext tests.

Updated KioskNextBackButtonTest so that TestApi is
recreated using the new ShelfView when a KioskNextUser logs in.

Updated ApplistButtonTest to not keep a local pointer to app_list_button_.

Bug: 951212
Change-Id: I72d3c3a3470ff776f63d841b39d9218b2bc00c70
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1618060
Commit-Queue: Yilkal Abe <yilkal@chromium.org>
Reviewed-by: Xiyuan Xia <xiyuan@chromium.org>
Cr-Commit-Position: refs/heads/master@{#661088}

ash/shelf/app_list_button_unittest.cc

@@ -50,12 +50,13 @@ class AppListButtonTest : public AshTestBase {
   // AshTestBase:
   void SetUp() override {
     AshTestBase::SetUp();
-    app_list_button_ =
-        GetPrimaryShelf()->GetShelfViewForTesting()->GetAppListButton();
   }
 
   void SendGestureEvent(ui::GestureEvent* event) {
-    app_list_button_->OnGestureEvent(event);
+    GetPrimaryShelf()
+        ->GetShelfViewForTesting()
+        ->GetAppListButton()
+        ->OnGestureEvent(event);
   }
 
   void SendGestureEventToSecondaryDisplay(ui::GestureEvent* event) {
@@ -68,11 +69,11 @@ class AppListButtonTest : public AshTestBase {
         ->OnGestureEvent(event);
   }
 
-  const AppListButton* app_list_button() const { return app_list_button_; }
+  const AppListButton* app_list_button() const {
+    return GetPrimaryShelf()->GetShelfViewForTesting()->GetAppListButton();
+  }
 
  private:
-  AppListButton* app_list_button_ = nullptr;
-
   DISALLOW_COPY_AND_ASSIGN(AppListButtonTest);
 };
 

ash/shelf/back_button_unittest.cc

@@ -47,7 +47,7 @@ class BackButtonTest : public AshTestBase {
     base::RunLoop().RunUntilIdle();
   }
 
- private:
+ protected:
   std::unique_ptr<ShelfViewTestAPI> test_api_;
 
   DISALLOW_COPY_AND_ASSIGN(BackButtonTest);
@@ -136,6 +136,14 @@ class KioskNextBackButtonTest : public BackButtonTest {
     client_ = BindMockKioskNextShellClient();
   }
 
+  void SimulateKioskNextSession() {
+    LogInKioskNextUser(GetSessionControllerClient());
+
+    // Update test_api_ because its reference to ShelfView is outdated.
+    test_api_ = std::make_unique<ShelfViewTestAPI>(
+        GetPrimaryShelf()->GetShelfViewForTesting());
+  }
+
  private:
   base::test::ScopedFeatureList scoped_feature_list_;
   std::unique_ptr<MockKioskNextShellClient> client_;
@@ -144,7 +152,7 @@ class KioskNextBackButtonTest : public BackButtonTest {
 };
 
 TEST_F(KioskNextBackButtonTest, BackKeySequenceGenerated) {
-  LogInKioskNextUser(GetSessionControllerClient());
+  SimulateKioskNextSession();
 
   // Tablet mode should be enabled in Kiosk Next.
   ASSERT_TRUE(Shell::Get()