chromeos: Makes AshTestBase/Helper target mash when appropriate

ash/test/ash_test_base.cc

Index: ash/test/ash_test_base.cc
diff --git a/ash/test/ash_test_base.cc b/ash/test/ash_test_base.cc
index c3dde233aae96220911d463116d46a773d7e2588..76ca2b0bc261dc39cd1f8584c780e56f30a42456 100644
--- a/ash/test/ash_test_base.cc
+++ b/ash/test/ash_test_base.cc
@@ -10,6 +10,7 @@
 #include "ash/common/test/test_session_state_delegate.h"
 #include "ash/common/test/test_system_tray_delegate.h"
 #include "ash/common/wm/window_positioner.h"
+#include "ash/common/wm_lookup.h"
 #include "ash/common/wm_shell.h"
 #include "ash/common/wm_window.h"
 #include "ash/display/extended_mouse_warp_controller.h"
@@ -63,9 +64,9 @@ class AshEventGeneratorDelegate
       const gfx::Point& point_in_screen) const override {
     display::Screen* screen = display::Screen::GetScreen();
     display::Display display = screen->GetDisplayNearestPoint(point_in_screen);
-    return Shell::GetInstance()
-        ->window_tree_host_manager()
+    return WmShell::Get()
         ->GetRootWindowForDisplayId(display.id())
+        ->aura_window()
         ->GetHost();
   }
 
@@ -134,7 +135,8 @@ void AshTestBase::SetUp() {
   // Move the mouse cursor to far away so that native events doesn't
   // interfere test expectations.
   Shell::GetPrimaryRootWindow()->MoveCursorTo(gfx::Point(-1000, -1000));
-  Shell::GetInstance()->cursor_manager()->EnableMouseEvents();
+  if (!WmShell::Get()->IsRunningInMash())
+    Shell::GetInstance()->cursor_manager()->EnableMouseEvents();
 
   // Changing GestureConfiguration shouldn't make tests fail. These values
   // prevent unexpected events from being generated during tests. Such as
@@ -196,8 +198,13 @@ display::Display::Rotation AshTestBase::GetCurrentInternalDisplayRotation() {
 
 // static
 void AshTestBase::UpdateDisplay(const std::string& display_specs) {
-  display::test::DisplayManagerTestApi(Shell::GetInstance()->display_manager())
-      .UpdateDisplay(display_specs);
+  if (WmShell::Get()->IsRunningInMash()) {
+    ash_test_helper_->UpdateDisplayForMash(display_specs);
+  } else {
+    display::test::DisplayManagerTestApi(
+        Shell::GetInstance()->display_manager())
+        .UpdateDisplay(display_specs);
+  }
 }
 
 aura::Window* AshTestBase::CurrentContext() {
@@ -264,9 +271,10 @@ aura::Window* AshTestBase::CreateTestWindowInShellWithDelegateAndType(
   } else {
     display::Display display =
         display::Screen::GetScreen()->GetDisplayMatching(bounds);
-    aura::Window* root = Shell::GetInstance()
-                             ->window_tree_host_manager()
-                             ->GetRootWindowForDisplayId(display.id());
+    aura::Window* root =
+        WmLookup::Get()
+            ->GetRootWindowControllerWithDisplayId(display.id())
+            ->GetRootWindow();
     gfx::Point origin = bounds.origin();
     ::wm::ConvertPointFromScreen(root, &origin);
     window->SetBounds(gfx::Rect(origin, bounds.size()));
@@ -388,5 +396,9 @@ void AshTestBase::SwapPrimaryDisplay() {
       display_manager()->GetSecondaryDisplay().id());
 }
 
+display::Display AshTestBase::GetSecondaryDisplay() {
+  return ash_test_helper_->GetSecondaryDisplay();
+}
+
 }  // namespace test
 }  // namespace ash