Make OnWmDisplayRemoved() take a WindowTreeHost

ui/aura/test/aura_test_base.cc

 // Copyright (c) 2012 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 "ui/aura/test/aura_test_base.h"
 
 #include "ui/aura/client/window_parenting_client.h"
 #include "ui/aura/mus/property_utils.h"
 #include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/mus/window_tree_host_mus.h"
@@ skipping 152 matching lines @@
     const std::string& name,
     std::unique_ptr<std::vector<uint8_t>>* new_data) {
   return true;
 }
 
 Window* AuraTestBase::OnWmCreateTopLevelWindow(
     ui::mojom::WindowType window_type,
     std::map<std::string, std::vector<uint8_t>>* properties) {
   Window* window = new Window(nullptr);
   SetWindowType(window, window_type);
+  window->Init(ui::LAYER_NOT_DRAWN);
   return window;
 }
 
 void AuraTestBase::OnWmClientJankinessChanged(
     const std::set<Window*>& client_windows,
     bool janky) {}
 
 void AuraTestBase::OnWmNewDisplay(
     std::unique_ptr<WindowTreeHostMus> window_tree_host,
     const display::Display& display) {
   // Take ownership of the WindowTreeHost.
   window_tree_host_mus_ = std::move(window_tree_host);
 }
 
-void AuraTestBase::OnWmDisplayRemoved(Window* window) {}
+void AuraTestBase::OnWmDisplayRemoved(WindowTreeHostMus* window_tree_host) {
+  if (window_tree_host_mus_.get() == window_tree_host)
+    window_tree_host_mus_.reset();
+}
 
 void AuraTestBase::OnWmDisplayModified(const display::Display& display) {}
 
 ui::mojom::EventResult AuraTestBase::OnAccelerator(uint32_t id,
                                                    const ui::Event& event) {
   return ui::mojom::EventResult::HANDLED;
 }
 
 void AuraTestBase::OnWmPerformMoveLoop(
     Window* window,
@@ skipping 19 matching lines @@
 
 void AuraTestBaseWithType::SetUp() {
   DCHECK(!setup_called_);
   setup_called_ = true;
   ConfigureBackend(GetParam());
   AuraTestBase::SetUp();
 }
 
 }  // namespace test
 }  // namespace aura