Fix LayerTreeHostCacheBehaviorOnOutputSurfaceRecreated test

cc/trees/layer_tree_host_unittest.cc

Index: cc/trees/layer_tree_host_unittest.cc
diff --git a/cc/trees/layer_tree_host_unittest.cc b/cc/trees/layer_tree_host_unittest.cc
index 248d766a694f4ef936daa95909394ae99c556afe..39e7a517b6151d278e71638bd862e6cc340facfd 100644
--- a/cc/trees/layer_tree_host_unittest.cc
+++ b/cc/trees/layer_tree_host_unittest.cc
@@ -576,27 +576,38 @@ class LayerTreeHostCacheBehaviorOnOutputSurfaceRecreated
  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: