Make OnWmDisplayRemoved() take a WindowTreeHost

ui/aura/test/aura_test_base.cc

Index: ui/aura/test/aura_test_base.cc
diff --git a/ui/aura/test/aura_test_base.cc b/ui/aura/test/aura_test_base.cc
index 22d0442f73bf94749851eeaddd6b5a603f0a26ef..ca5d93f9c5009697d696d00e885354728a339cb9 100644
--- a/ui/aura/test/aura_test_base.cc
+++ b/ui/aura/test/aura_test_base.cc
@@ -170,6 +170,7 @@ Window* AuraTestBase::OnWmCreateTopLevelWindow(
     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;
 }
 
@@ -184,7 +185,10 @@ void AuraTestBase::OnWmNewDisplay(
   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) {}