mus: Fix ash::GetRootWindowController for child windows/widgets

ash/root_window_controller_unittest.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 "ash/root_window_controller.h"
 
 #include <memory>
 
 #include "ash/display/display_manager.h"
 #include "ash/session/session_state_delegate.h"
@@ skipping 556 matching lines @@
   EXPECT_EQ(NULL, controllers[1]->GetWindowForFullscreenMode());
 
   // Verify that the first root window controller remains in fullscreen mode
   // when a window on the other display is activated.
   w3->Activate();
   EXPECT_EQ(w2->GetNativeWindow(),
             controllers[0]->GetWindowForFullscreenMode());
   EXPECT_EQ(NULL, controllers[1]->GetWindowForFullscreenMode());
 }
 
+// Test that GetRootWindowController() works with multiple displays and
+// child widgets.
+TEST_F(RootWindowControllerTest, GetRootWindowController) {
+  if (!SupportsMultipleDisplays())
+    return;
+  UpdateDisplay("600x600,600x600");
+  Shell::RootWindowControllerList controllers =
+      Shell::GetInstance()->GetAllRootWindowControllers();
+  ASSERT_EQ(2u, controllers.size());
+
+  // Test null.
+  EXPECT_FALSE(GetRootWindowController(nullptr));
+
+  // Test a widget on the first display.
+  Widget* w1 = CreateTestWidget(gfx::Rect(0, 0, 100, 100));
+  EXPECT_EQ(controllers[0],
+            GetRootWindowController(w1->GetNativeWindow()->GetRootWindow()));
+
+  // Test a child widget.
+  Widget* w2 = Widget::CreateWindowWithParentAndBounds(
+      nullptr, w1->GetNativeWindow(), gfx::Rect(0, 0, 100, 100));
+  EXPECT_EQ(controllers[0],
+            GetRootWindowController(w2->GetNativeWindow()->GetRootWindow()));
+
+  // Test a widget on the second display.
+  Widget* w3 = CreateTestWidget(gfx::Rect(600, 0, 100, 100));
+  EXPECT_EQ(controllers[1],
+            GetRootWindowController(w3->GetNativeWindow()->GetRootWindow()));
+}
+
 // Test that user session window can't be focused if user session blocked by
 // some overlapping UI.
 TEST_F(RootWindowControllerTest, FocusBlockedWindow) {
   UpdateDisplay("600x600");
   RootWindowController* controller =
       Shell::GetInstance()->GetPrimaryRootWindowController();
   aura::Window* lock_container =
       controller->GetContainer(kShellWindowId_LockScreenContainer);
   aura::Window* lock_window = Widget::CreateWindowWithParentAndBounds(NULL,
       lock_container, gfx::Rect(0, 0, 100, 100))->GetNativeView();
@@ skipping 432 matching lines @@
   keyboard_controller->ui()->GetKeyboardWindow()->SetBounds(
       gfx::Rect(0, 400, 800, 200));
   EXPECT_EQ("0,400 800x200", keyboard_container->bounds().ToString());
 
   UpdateDisplay("600x800");
   EXPECT_EQ("0,600 600x200", keyboard_container->bounds().ToString());
 }
 
 }  // namespace test
 }  // namespace ash