Adds support for multiple displays to WmTestBase

ash/mus/window_manager.h

 // Copyright 2015 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.
 
 #ifndef ASH_MUS_WINDOW_MANAGER_H_
 #define ASH_MUS_WINDOW_MANAGER_H_
 
 #include <stdint.h>
 
 #include <memory>
@@ skipping 76 matching lines @@
   void AddAcceleratorHandler(uint16_t id_namespace,
                              AcceleratorHandler* handler);
   void RemoveAcceleratorHandler(uint16_t id_namespace);
 
   void AddObserver(WindowManagerObserver* observer);
   void RemoveObserver(WindowManagerObserver* observer);
 
  private:
   friend class WmTestHelper;
 
+  using RootWindowControllers = std::set<std::unique_ptr<RootWindowController>>;
+
   RootWindowController* CreateRootWindowController(
       ui::Window* window,
       const display::Display& display);
 
+  // Deletes the specified RootWindowController. Called when a display is
+  // removed.
+  void DestroyRootWindowController(
+      RootWindowController* root_window_controller);
+
   void Shutdown();
 
+  // Returns an iterator into |root_window_controllers_|. Returns
+  // root_window_controllers_.end() if |window| is not the root of a
+  // RootWindowController.
+  RootWindowControllers::iterator FindRootWindowControllerByWindow(
+      ui::Window* window);
+
   // ui::WindowObserver:
   void OnWindowDestroying(ui::Window* window) override;
   void OnWindowDestroyed(ui::Window* window) override;
 
   // WindowTreeClientDelegate:
   void OnEmbed(ui::Window* root) override;
   void OnEmbedRootDestroyed(ui::Window* root) override;
   void OnLostConnection(ui::WindowTreeClient* client) override;
   void OnPointerEventObserved(const ui::PointerEvent& event,
                               ui::Window* target) override;
@@ skipping 23 matching lines @@
 
   std::unique_ptr<ui::WindowTreeClient> window_tree_client_;
 
   ui::WindowManagerClient* window_manager_client_ = nullptr;
 
   std::unique_ptr<views::PointerWatcherEventRouter>
       pointer_watcher_event_router_;
 
   std::unique_ptr<ShadowController> shadow_controller_;
 
-  std::set<std::unique_ptr<RootWindowController>> root_window_controllers_;
+  RootWindowControllers root_window_controllers_;
 
   base::ObserverList<WindowManagerObserver> observers_;
 
   std::unique_ptr<display::ScreenBase> screen_;
 
   std::unique_ptr<WmShellMus> shell_;
 
   std::unique_ptr<WmLookupMus> lookup_;
 
   std::map<uint16_t, AcceleratorHandler*> accelerator_handlers_;
   uint16_t next_accelerator_namespace_id_ = 0u;
 
   DISALLOW_COPY_AND_ASSIGN(WindowManager);
 };
 
 }  // namespace mus
 }  // namespace ash
 
 #endif  // ASH_MUS_WINDOW_MANAGER_H_