Fix LayerTreeHostCacheBehaviorOnOutputSurfaceRecreated test

cc/trees/layer_tree_host_unittest.cc

 // Copyright 2011 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "cc/trees/layer_tree_host.h"
 
 #include <stddef.h>
 #include <stdint.h>
 
 #include <algorithm>
@@ skipping 558 matching lines @@
 
 SINGLE_AND_MULTI_THREAD_TEST_F(LayerTreeHostFreeContextResourcesOnDestroy);
 
 // Test if the LTH successfully frees and stops freeing context resources
 // when the OutputSurface is lost and recreated.
 class LayerTreeHostCacheBehaviorOnOutputSurfaceRecreated
     : public LayerTreeHostContextCacheTest {
  public:
   void WillBeginImplFrameOnThread(LayerTreeHostImpl* host_impl,
                                   const BeginFrameArgs& args) override {
+    // This code is run once, to trigger recreation of our OutputSurface.
+    if (has_recreated_)
+      return;
+
     // Ensure that our initialization expectations have completed.
     Mock::VerifyAndClearExpectations(mock_main_context_support_);
     Mock::VerifyAndClearExpectations(mock_worker_context_support_);
 
-    if (has_recreated_) {
-      // Destruction exptectations.
-      EXPECT_CALL(*mock_worker_context_support_,
-                  SetAggressivelyFreeResources(true));
-      EXPECT_CALL(*mock_main_context_support_,
-                  SetAggressivelyFreeResources(true));
-      EndTest();
-      return;
-    }
-    has_recreated_ = true;
-
     // Output surface lost expectations.
     EXPECT_CALL(*mock_worker_context_support_,
                 SetAggressivelyFreeResources(true));
     EXPECT_CALL(*mock_main_context_support_,
                 SetAggressivelyFreeResources(true));
     host_impl->DidLoseOutputSurface();
+    has_recreated_ = true;
+  }
+
+  void DidInitializeOutputSurface() override {
+    // This is run after we have recreated our OutputSurface.
+    if (!has_recreated_)
+      return;
+
+    // Ensure that our initialization expectations have completed.
+    Mock::VerifyAndClearExpectations(mock_main_context_support_);
+    Mock::VerifyAndClearExpectations(mock_worker_context_support_);
+
+    // Destruction exptectations.
+    EXPECT_CALL(*mock_worker_context_support_,
+                SetAggressivelyFreeResources(true));
+    EXPECT_CALL(*mock_main_context_support_,
+                SetAggressivelyFreeResources(true));
+    EndTest();
   }
 
  private:
   bool has_recreated_ = false;
 };
 
 SINGLE_AND_MULTI_THREAD_TEST_F(
     LayerTreeHostCacheBehaviorOnOutputSurfaceRecreated);
 
 // Two setNeedsCommits in a row should lead to at least 1 commit and at least 1
@@ skipping 6415 matching lines @@
  private:
   FakeContentLayerClient client_;
   const gfx::Size viewport_size_;
   const gfx::Size large_image_size_;
 };
 
 SINGLE_AND_MULTI_THREAD_TEST_F(GpuRasterizationSucceedsWithLargeImage);
 
 }  // namespace
 }  // namespace cc