Add Screen::GetDisplayForNewWindows to support remote mojo apps

Chrome OS supports the concept of a "root window for new windows",
which allows new windows to be automatically placed on the correct
display. This knowledge lives in ash, currently inside the browser
process.

Remote processes that use views::Widget (e.g. the shortcut viewer app)
sometimes need to know what window they are going to open on,
for example to do initial bounds computations.

Change Screen to provide set and get for the display for new windows.
Use mojo to send the information to other processes via ScreenProvider
(on the ash side) and ScreenMus (on the remote side). Send the display
id as part of display updates so that ScreenMus gets it on startup,
and also so updates are atomic with other display changes.

Once this lands we'll be able to port browser windows to use this
mechanism and delete the ash::ShellState / chrome ShellStateClient
code I built as a one-off for browser windows under mash.

Bug: 867559, 764009, 826569
Test: added unit tests
Change-Id: I6b5beb2610243f30072c3d73180b3ce21a310ca6
Reviewed-on: https://chromium-review.googlesource.com/1162980Reviewed-by: Scott Violet <sky@chromium.org>
Reviewed-by: Tom Sepez <tsepez@chromium.org>
Reviewed-by: Michael Wasserman <msw@chromium.org>
Commit-Queue: James Cook <jamescook@chromium.org>
Cr-Commit-Position: refs/heads/master@{#581055}

ash/display/screen_ash.cc

@@ -10,6 +10,7 @@
 #include "ash/root_window_settings.h"
 #include "ash/shelf/shelf_widget.h"
 #include "ash/shell.h"
+#include "ash/shell_state.h"
 #include "ash/wm/root_window_finder.h"
 #include "base/command_line.h"
 #include "base/logging.h"
@@ -68,6 +69,9 @@ class ScreenForShutdown : public display::Screen {
   display::Display GetPrimaryDisplay() const override {
     return primary_display_;
   }
+  display::Display GetDisplayForNewWindows() const override {
+    return primary_display_;
+  }
   void AddObserver(display::DisplayObserver* observer) override {
     NOTREACHED() << "Observer should not be added during shutdown";
   }
@@ -168,6 +172,11 @@ display::Display ScreenAsh::GetPrimaryDisplay() const {
       WindowTreeHostManager::GetPrimaryDisplayId());
 }
 
+display::Display ScreenAsh::GetDisplayForNewWindows() const {
+  return GetDisplayNearestWindow(
+      Shell::Get()->shell_state()->GetRootWindowForNewWindows());
+}
+
 void ScreenAsh::AddObserver(display::DisplayObserver* observer) {
   GetDisplayManager()->AddObserver(observer);
 }

ash/display/screen_ash.h

@@ -42,6 +42,7 @@ class ASH_EXPORT ScreenAsh : public display::Screen {
   display::Display GetDisplayMatching(
       const gfx::Rect& match_rect) const override;
   display::Display GetPrimaryDisplay() const override;
+  display::Display GetDisplayForNewWindows() const override;
   void AddObserver(display::DisplayObserver* observer) override;
   void RemoveObserver(display::DisplayObserver* observer) override;
 

ash/display/screen_ash_unittest.cc

@@ -2,6 +2,7 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
+#include "ash/scoped_root_window_for_new_windows.h"
 #include "ash/shell.h"
 #include "ash/test/ash_test_base.h"
 #include "ui/aura/test/test_window_delegate.h"
@@ -31,4 +32,19 @@ TEST_F(ScreenAshTest, TestGetWindowAtScreenPoint) {
                             gfx::Point(250, 260)));
 }
 
+TEST_F(ScreenAshTest, GetDisplayForNewWindows) {
+  UpdateDisplay("200x200,400x400");
+  display::Screen* screen = display::Screen::GetScreen();
+  const std::vector<display::Display> displays = screen->GetAllDisplays();
+  ASSERT_EQ(2u, displays.size());
+
+  // The display for new windows defaults to primary display.
+  EXPECT_EQ(displays[0].id(), screen->GetDisplayForNewWindows().id());
+
+  // The display for new windows is updated when the root window for new windows
+  // changes.
+  ScopedRootWindowForNewWindows scoped_root(Shell::GetAllRootWindows()[1]);
+  EXPECT_EQ(displays[1].id(), screen->GetDisplayForNewWindows().id());
+}
+
 }  // namespace ash

ash/shell.cc

@@ -855,9 +855,6 @@ Shell::~Shell() {
   display_speaker_controller_.reset();
   screen_switch_check_controller_.reset();
 
-  // This also deletes all RootWindows. Note that we invoke Shutdown() on
-  // WindowTreeHostManager before resetting |window_tree_host_manager_|, since
-  // destruction of its owned RootWindowControllers relies on the value.
   ScreenAsh::CreateScreenForShutdown();
   display_configuration_controller_.reset();
 
@@ -891,11 +888,17 @@ Shell::~Shell() {
   ime_focus_handler_.reset();
 
   shell_port_->Shutdown();
+
+  // Stop observing window activation changes before closing all windows.
+  focus_controller_->RemoveObserver(this);
+
+  // This also deletes all RootWindows. Note that we invoke Shutdown() on
+  // WindowTreeHostManager before resetting |window_tree_host_manager_|, since
+  // destruction of its owned RootWindowControllers relies on the value.
   window_tree_host_manager_->Shutdown();
 
   // Depends on |focus_controller_|, so must be destroyed before.
   window_tree_host_manager_.reset();
-  focus_controller_->RemoveObserver(this);
   focus_controller_.reset();
   screen_position_controller_.reset();
 

ash/shell_state.cc

@@ -7,6 +7,9 @@
 #include <memory>
 #include <utility>
 
+#include "ash/shell.h"
+#include "ash/ws/window_service_owner.h"
+#include "services/ui/ws2/window_service.h"
 #include "ui/display/display.h"
 #include "ui/display/screen.h"
 
@@ -48,6 +51,14 @@ void ShellState::NotifyAllClients() {
   clients_.ForAllPtrs([display_id](mojom::ShellStateClient* client) {
     client->SetDisplayIdForNewWindows(display_id);
   });
+
+  // WindowService broadcasts the display id over mojo to all remote apps.
+  // TODO(jamescook): Move this into Shell when ShellState is removed.
+  WindowServiceOwner* ws_owner = Shell::Get()->window_service_owner();
+  // |ws_owner| is null during shutdown and tests. |window_service()| is null
+  // during early startup.
+  if (ws_owner && ws_owner->window_service())
+    ws_owner->window_service()->SetDisplayForNewWindows(display_id);
 }
 
 int64_t ShellState::GetDisplayIdForNewWindows() const {

ash/shell_state.h

@@ -22,6 +22,9 @@ class Window;
 namespace ash {
 
 // Provides access via mojo to ash::Shell state.
+// TODO(jamescook): Move |root_window_for_new_windows_| to Shell, convert
+// browser code to use display::Screen::GetDisplayForNewWindows() and delete
+// this class.
 class ASH_EXPORT ShellState : public mojom::ShellState {
  public:
   ShellState();

ash/ws/window_service_owner.cc

@@ -13,9 +13,12 @@
 #include "services/ui/ws2/gpu_interface_provider.h"
 #include "services/ui/ws2/window_service.h"
 #include "ui/base/ui_base_features.h"
+#include "ui/display/display.h"
+#include "ui/display/screen.h"
 #include "ui/wm/core/focus_controller.h"
 
 namespace ash {
+
 WindowServiceOwner::WindowServiceOwner(
     std::unique_ptr<ui::ws2::GpuInterfaceProvider> gpu_interface_provider)
     : gpu_interface_provider_(std::move(gpu_interface_provider)) {}
@@ -38,6 +41,8 @@ void WindowServiceOwner::BindWindowService(
   window_service_->SetFrameDecorationValues(
       NonClientFrameController::GetPreferredClientAreaInsets(),
       NonClientFrameController::GetMaxTitleBarButtonWidth());
+  window_service_->SetDisplayForNewWindows(
+      display::Screen::GetScreen()->GetDisplayForNewWindows().id());
   RegisterWindowProperties(window_service_->property_converter());
   service_context_ = std::make_unique<service_manager::ServiceContext>(
       std::move(window_service), std::move(request));

services/ui/public/interfaces/screen_provider_observer.mojom

@@ -12,8 +12,11 @@ interface ScreenProviderObserver {
   // that WS is running on has an integrated display, for example a laptop
   // internal display, then |internal_display_id| will be the corresponding
   // dislay id. If there is no internal display then |internal_display_id| will
-  // be kInvalidDisplayID.
+  // be kInvalidDisplayID. |display_id_for_new_windows| is the display on which
+  // to place new top-level windows, usually the display on which a window was
+  // last activated.
   OnDisplaysChanged(array<WsDisplay> displays,
                     int64 primary_display_id,
-                    int64 internal_display_id);
+                    int64 internal_display_id,
+                    int64 display_id_for_new_windows);
 };

services/ui/ws2/screen_provider.cc

@@ -49,6 +49,13 @@ void ScreenProvider::SetFrameDecorationValues(
   max_title_bar_button_width_ = max_title_bar_button_width;
 }
 
+void ScreenProvider::SetDisplayForNewWindows(int64_t display_id) {
+  if (display_id == display_id_for_new_windows_)
+    return;
+  display_id_for_new_windows_ = display_id;
+  NotifyAllObservers();
+}
+
 void ScreenProvider::DisplayMetricsChanged(const display::Display& display,
                                            uint32_t changed_metrics) {
   if ((changed_metrics &
@@ -70,7 +77,8 @@ void ScreenProvider::NotifyAllObservers() {
 
 void ScreenProvider::NotifyObserver(mojom::ScreenProviderObserver* observer) {
   observer->OnDisplaysChanged(GetAllDisplays(), GetPrimaryDisplayId(),
-                              GetInternalDisplayId());
+                              GetInternalDisplayId(),
+                              display_id_for_new_windows_);
 }
 
 std::vector<mojom::WsDisplayPtr> ScreenProvider::GetAllDisplays() {

services/ui/ws2/screen_provider.h

@@ -13,6 +13,7 @@
 #include "services/ui/public/interfaces/screen_provider_observer.mojom.h"
 #include "services/ui/public/interfaces/window_tree_constants.mojom.h"
 #include "ui/display/display_observer.h"
+#include "ui/display/types/display_constants.h"
 #include "ui/gfx/geometry/insets.h"
 
 namespace display {
@@ -37,6 +38,9 @@ class COMPONENT_EXPORT(WINDOW_SERVICE) ScreenProvider
   void SetFrameDecorationValues(const gfx::Insets& client_area_insets,
                                 int max_title_bar_button_width);
 
+  // See WindowService documentation.
+  void SetDisplayForNewWindows(int64_t display_id);
+
   // See comment in WindowService as to why this is special cased.
   void DisplayMetricsChanged(const display::Display& display,
                              uint32_t changed_metrics);
@@ -58,6 +62,8 @@ class COMPONENT_EXPORT(WINDOW_SERVICE) ScreenProvider
   gfx::Insets client_area_insets_;
   int max_title_bar_button_width_ = 0;
 
+  int64_t display_id_for_new_windows_ = display::kInvalidDisplayId;
+
   base::ObserverList<mojom::ScreenProviderObserver> observers_;
 
   DISALLOW_COPY_AND_ASSIGN(ScreenProvider);

services/ui/ws2/screen_provider_unittest.cc

@@ -93,6 +93,25 @@ TEST(ScreenProviderTest, AddRemoveDisplay) {
   EXPECT_EQ(222, observer.primary_display_id());
 }
 
+TEST(ScreenProviderTest, SetDisplayForNewWindows) {
+  // Set up 2 displays.
+  constexpr int64_t kDisplayId1 = 111;
+  constexpr int64_t kDisplayId2 = 222;
+  TestScreen screen;
+  screen.AddDisplay(Display(kDisplayId1), DisplayList::Type::PRIMARY);
+  screen.AddDisplay(Display(kDisplayId2), DisplayList::Type::NOT_PRIMARY);
+
+  // Set the display for new windows to the second display.
+  ScreenProvider screen_provider;
+  screen_provider.SetDisplayForNewWindows(kDisplayId2);
+
+  TestScreenProviderObserver observer;
+  screen_provider.AddObserver(&observer);
+
+  // The screen information includes the display for new windows.
+  EXPECT_EQ(kDisplayId2, observer.display_id_for_new_windows());
+}
+
 TEST(ScreenProviderTest, SetFrameDecorationValues) {
   // Set up a single display.
   TestScreen screen;

services/ui/ws2/test_screen_provider_observer.cc

@@ -30,11 +30,13 @@ TestScreenProviderObserver::~TestScreenProviderObserver() = default;
 void TestScreenProviderObserver::OnDisplaysChanged(
     std::vector<mojom::WsDisplayPtr> displays,
     int64_t primary_display_id,
-    int64_t internal_display_id) {
+    int64_t internal_display_id,
+    int64_t display_id_for_new_windows) {
   displays_ = std::move(displays);
   display_ids_ = DisplayIdsToString(displays_);
   primary_display_id_ = primary_display_id;
   internal_display_id_ = internal_display_id;
+  display_id_for_new_windows_ = display_id_for_new_windows;
 }
 
 }  // namespace ws2

services/ui/ws2/test_screen_provider_observer.h

@@ -13,6 +13,7 @@
 #include "base/macros.h"
 #include "services/ui/public/interfaces/screen_provider_observer.mojom.h"
 #include "services/ui/ws2/window_service_delegate.h"
+#include "ui/display/types/display_constants.h"
 #include "ui/events/event.h"
 
 namespace ui {
@@ -29,17 +30,22 @@ class TestScreenProviderObserver : public mojom::ScreenProviderObserver {
   std::string& display_ids() { return display_ids_; }
   int64_t primary_display_id() const { return primary_display_id_; }
   int64_t internal_display_id() const { return internal_display_id_; }
+  int64_t display_id_for_new_windows() const {
+    return display_id_for_new_windows_;
+  }
 
   // mojom::ScreenProviderObserver:
   void OnDisplaysChanged(std::vector<mojom::WsDisplayPtr> displays,
                          int64_t primary_display_id,
-                         int64_t internal_display_id) override;
+                         int64_t internal_display_id,
+                         int64_t display_id_for_new_windows) override;
 
  private:
   std::vector<mojom::WsDisplayPtr> displays_;
   std::string display_ids_;
-  int64_t primary_display_id_ = 0;
-  int64_t internal_display_id_ = 0;
+  int64_t primary_display_id_ = display::kInvalidDisplayId;
+  int64_t internal_display_id_ = display::kInvalidDisplayId;
+  int64_t display_id_for_new_windows_ = display::kInvalidDisplayId;
 
   DISALLOW_COPY_AND_ASSIGN(TestScreenProviderObserver);
 };

services/ui/ws2/window_service.cc

@@ -100,6 +100,10 @@ void WindowService::SetFrameDecorationValues(
                                              max_title_bar_button_width);
 }
 
+void WindowService::SetDisplayForNewWindows(int64_t display_id) {
+  screen_provider_->SetDisplayForNewWindows(display_id);
+}
+
 // static
 bool WindowService::HasRemoteClient(const aura::Window* window) {
   return ServerWindow::GetMayBeNull(window);

services/ui/ws2/window_service.h

@@ -90,6 +90,10 @@ class COMPONENT_EXPORT(WINDOW_SERVICE) WindowService
   void SetFrameDecorationValues(const gfx::Insets& client_area_insets,
                                 int max_title_bar_button_width);
 
+  // Sets the display to use for new windows. The window server broadcasts this
+  // over mojo to all remote clients.
+  void SetDisplayForNewWindows(int64_t display_id);
+
   // Whether |window| hosts a remote client.
   static bool HasRemoteClient(const aura::Window* window);
 

ui/aura/mus/window_tree_client.cc

@@ -1511,10 +1511,11 @@ void WindowTreeClient::GetScreenProviderObserver(
 void WindowTreeClient::OnDisplaysChanged(
     std::vector<ui::mojom::WsDisplayPtr> ws_displays,
     int64_t primary_display_id,
-    int64_t internal_display_id) {
+    int64_t internal_display_id,
+    int64_t display_id_for_new_windows) {
   got_initial_displays_ = true;
   delegate_->OnDisplaysChanged(std::move(ws_displays), primary_display_id,
-                               internal_display_id);
+                               internal_display_id, display_id_for_new_windows);
 }
 
 void WindowTreeClient::RequestClose(ui::Id window_id) {

ui/aura/mus/window_tree_client.h

@@ -438,7 +438,8 @@ class AURA_EXPORT WindowTreeClient
   // ui::mojom::ScreenProviderObserver:
   void OnDisplaysChanged(std::vector<ui::mojom::WsDisplayPtr> ws_displays,
                          int64_t primary_display_id,
-                         int64_t internal_display_id) override;
+                         int64_t internal_display_id,
+                         int64_t display_id_for_new_windows) override;
 
   // Overriden from WindowTreeHostMusDelegate:
   void OnWindowTreeHostBoundsWillChange(

ui/aura/mus/window_tree_client_delegate.h

@@ -70,7 +70,8 @@ class AURA_EXPORT WindowTreeClientDelegate {
   virtual void OnDisplaysChanged(
       std::vector<ui::mojom::WsDisplayPtr> ws_displays,
       int64_t primary_display_id,
-      int64_t internal_display_id) {}
+      int64_t internal_display_id,
+      int64_t display_id_for_new_windows) {}
 
  protected:
   virtual ~WindowTreeClientDelegate() {}

ui/display/screen.cc

@@ -38,6 +38,10 @@ Display Screen::GetDisplayNearestView(gfx::NativeView view) const {
   return GetDisplayNearestWindow(GetWindowForView(view));
 }
 
+Display Screen::GetDisplayForNewWindows() const {
+  return GetPrimaryDisplay();
+}
+
 gfx::Rect Screen::ScreenToDIPRectInWindow(gfx::NativeView view,
                                           const gfx::Rect& screen_rect) const {
   float scale = GetDisplayNearestView(view).device_scale_factor();

ui/display/screen.h

@@ -74,6 +74,10 @@ class DISPLAY_EXPORT Screen {
   // Returns the primary display.
   virtual Display GetPrimaryDisplay() const = 0;
 
+  // Returns a suggested display to use when creating a new window. On most
+  // platforms just returns the primary display.
+  virtual Display GetDisplayForNewWindows() const;
+
   // Adds/Removes display observers.
   virtual void AddObserver(DisplayObserver* observer) = 0;
   virtual void RemoveObserver(DisplayObserver* observer) = 0;

ui/display/screen_unittest.cc

@@ -58,4 +58,12 @@ TEST_F(ScreenTest, GetDisplayWithDisplayId) {
   EXPECT_EQ(DEFAULT_DISPLAY_HEIGHT, display.size().height());
 }
 
+TEST_F(ScreenTest, GetDisplayForNewWindows) {
+  Screen* screen = Screen::GetScreen();
+
+  // Display for new windows defaults to the primary display.
+  EXPECT_EQ(screen->GetPrimaryDisplay().id(),
+            screen->GetDisplayForNewWindows().id());
+}
+
 }  // namespace display

ui/views/mus/mus_client.cc

@@ -372,9 +372,10 @@ void MusClient::OnPointerEventObserved(const ui::PointerEvent& event,
 void MusClient::OnDisplaysChanged(
     std::vector<ui::mojom::WsDisplayPtr> ws_displays,
     int64_t primary_display_id,
-    int64_t internal_display_id) {
+    int64_t internal_display_id,
+    int64_t display_id_for_new_windows) {
   screen_->OnDisplaysChanged(std::move(ws_displays), primary_display_id,
-                             internal_display_id);
+                             internal_display_id, display_id_for_new_windows);
 }
 
 void MusClient::OnWindowManagerFrameValuesChanged() {

ui/views/mus/mus_client.h

@@ -168,7 +168,8 @@ class VIEWS_MUS_EXPORT MusClient : public aura::WindowTreeClientDelegate,
   aura::PropertyConverter* GetPropertyConverter() override;
   void OnDisplaysChanged(std::vector<ui::mojom::WsDisplayPtr> ws_displays,
                          int64_t primary_display_id,
-                         int64_t internal_display_id) override;
+                         int64_t internal_display_id,
+                         int64_t display_id_for_new_windows) override;
 
   // ScreenMusDelegate:
   void OnWindowManagerFrameValuesChanged() override;

ui/views/mus/screen_mus.cc

@@ -47,7 +47,8 @@ ScreenMus::~ScreenMus() {
 void ScreenMus::OnDisplaysChanged(
     std::vector<ui::mojom::WsDisplayPtr> ws_displays,
     int64_t primary_display_id,
-    int64_t internal_display_id) {
+    int64_t internal_display_id,
+    int64_t display_id_for_new_windows) {
   const bool primary_changed = primary_display_id != GetPrimaryDisplay().id();
   int64_t handled_display_id = display::kInvalidDisplayId;
   const WindowManagerFrameValues initial_frame_values =
@@ -102,6 +103,9 @@ void ScreenMus::OnDisplaysChanged(
       initial_frame_values != WindowManagerFrameValues::instance()) {
     delegate_->OnWindowManagerFrameValuesChanged();
   }
+
+  // GetDisplayForNewWindows() can handle ids that are not in the list.
+  display_id_for_new_windows_ = display_id_for_new_windows;
 }
 
 display::Display ScreenMus::GetDisplayNearestWindow(
@@ -126,4 +130,13 @@ aura::Window* ScreenMus::GetWindowAtScreenPoint(const gfx::Point& point) {
   return delegate_->GetWindowAtScreenPoint(point);
 }
 
+display::Display ScreenMus::GetDisplayForNewWindows() const {
+  display::Display display;
+  if (GetDisplayWithDisplayId(display_id_for_new_windows_, &display))
+    return display;
+
+  // Fallback to primary display.
+  return GetPrimaryDisplay();
+}
+
 }  // namespace views

ui/views/mus/screen_mus.h

@@ -24,10 +24,8 @@ class VIEWS_MUS_EXPORT ScreenMus : public display::ScreenBase,
   // ui::mojom::ScreenProviderObserver:
   void OnDisplaysChanged(std::vector<ui::mojom::WsDisplayPtr> ws_displays,
                          int64_t primary_display_id,
-                         int64_t internal_display_id) override;
-
- private:
-  friend class ScreenMusTestApi;
+                         int64_t internal_display_id,
+                         int64_t display_id_for_new_windows) override;
 
   // display::Screen:
   display::Display GetDisplayNearestWindow(
@@ -35,8 +33,11 @@ class VIEWS_MUS_EXPORT ScreenMus : public display::ScreenBase,
   gfx::Point GetCursorScreenPoint() override;
   bool IsWindowUnderCursor(gfx::NativeWindow window) override;
   aura::Window* GetWindowAtScreenPoint(const gfx::Point& point) override;
+  display::Display GetDisplayForNewWindows() const override;
 
+ private:
   ScreenMusDelegate* delegate_;
+  int64_t display_id_for_new_windows_ = display::kInvalidDisplayId;
 
   DISALLOW_COPY_AND_ASSIGN(ScreenMus);
 };

ui/views/mus/screen_mus_unittest.cc

@@ -6,28 +6,12 @@
 
 #include "base/command_line.h"
 #include "base/test/scoped_task_environment.h"
+#include "testing/gtest/include/gtest/gtest.h"
 #include "ui/display/display_switches.h"
 #include "ui/display/screen.h"
 #include "ui/views/test/scoped_views_test_helper.h"
-#include "ui/views/test/views_test_base.h"
 
 namespace views {
-
-class ScreenMusTestApi {
- public:
-  static void CallOnDisplaysChanged(
-      ScreenMus* screen,
-      std::vector<ui::mojom::WsDisplayPtr> ws_displays,
-      int64_t primary_display_id,
-      int64_t internal_display_id) {
-    screen->OnDisplaysChanged(std::move(ws_displays), primary_display_id,
-                              internal_display_id);
-  }
-
- private:
-  DISALLOW_IMPLICIT_CONSTRUCTORS(ScreenMusTestApi);
-};
-
 namespace {
 
 std::vector<ui::mojom::WsDisplayPtr> ConvertDisplayToWsDisplays(
@@ -43,9 +27,10 @@ std::vector<ui::mojom::WsDisplayPtr> ConvertDisplayToWsDisplays(
   return results;
 }
 
-TEST(ScreenMusTest, ConsistentDisplayInHighDPI) {
+TEST(ScreenMusScaleFactorTest, ConsistentDisplayInHighDPI) {
   base::test::ScopedTaskEnvironment task_environment(
       base::test::ScopedTaskEnvironment::MainThreadType::UI);
+  // Must be set before |test_helper| is constructed.
   base::CommandLine::ForCurrentProcess()->AppendSwitchASCII(
       switches::kForceDeviceScaleFactor, "2");
   ScopedViewsTestHelper test_helper;
@@ -58,10 +43,20 @@ TEST(ScreenMusTest, ConsistentDisplayInHighDPI) {
   }
 }
 
-TEST(ScreenMusTest, PrimaryChangedToExisting) {
-  base::test::ScopedTaskEnvironment task_environment(
-      base::test::ScopedTaskEnvironment::MainThreadType::UI);
-  ScopedViewsTestHelper test_helper;
+class ScreenMusTest : public testing::Test {
+ public:
+  ScreenMusTest() = default;
+  ~ScreenMusTest() override = default;
+
+ private:
+  base::test::ScopedTaskEnvironment task_environment_{
+      base::test::ScopedTaskEnvironment::MainThreadType::UI};
+  ScopedViewsTestHelper test_helper_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScreenMusTest);
+};
+
+TEST_F(ScreenMusTest, PrimaryChangedToExisting) {
   ScreenMus* screen = static_cast<ScreenMus*>(display::Screen::GetScreen());
   std::vector<display::Display> displays = screen->GetAllDisplays();
   ASSERT_FALSE(displays.empty());
@@ -69,17 +64,14 @@ TEST(ScreenMusTest, PrimaryChangedToExisting) {
   // Convert to a single display with a different primary id.
   displays.resize(1);
   displays[0].set_id(displays[0].id() + 1);
-  ScreenMusTestApi::CallOnDisplaysChanged(
-      screen, ConvertDisplayToWsDisplays(displays), displays[0].id(), 0);
+  screen->OnDisplaysChanged(ConvertDisplayToWsDisplays(displays),
+                            displays[0].id(), 0, 0);
   ASSERT_EQ(1u, screen->GetAllDisplays().size());
   EXPECT_EQ(displays[0].id(), screen->GetAllDisplays()[0].id());
   EXPECT_EQ(displays[0].id(), screen->GetPrimaryDisplay().id());
 }
 
-TEST(ScreenMusTest, AddAndUpdate) {
-  base::test::ScopedTaskEnvironment task_environment(
-      base::test::ScopedTaskEnvironment::MainThreadType::UI);
-  ScopedViewsTestHelper test_helper;
+TEST_F(ScreenMusTest, AddAndUpdate) {
   ScreenMus* screen = static_cast<ScreenMus*>(display::Screen::GetScreen());
   std::vector<display::Display> displays = screen->GetAllDisplays();
   ASSERT_FALSE(displays.empty());
@@ -91,8 +83,8 @@ TEST(ScreenMusTest, AddAndUpdate) {
   displays[0].set_bounds(new_bounds);
   displays.push_back(displays[0]);
   displays[1].set_id(displays[0].id() + 1);
-  ScreenMusTestApi::CallOnDisplaysChanged(
-      screen, ConvertDisplayToWsDisplays(displays), displays[1].id(), 0);
+  screen->OnDisplaysChanged(ConvertDisplayToWsDisplays(displays),
+                            displays[1].id(), 0, 0);
   ASSERT_EQ(2u, screen->GetAllDisplays().size());
   ASSERT_TRUE(screen->display_list().FindDisplayById(displays[0].id()) !=
               screen->display_list().displays().end());
@@ -105,5 +97,31 @@ TEST(ScreenMusTest, AddAndUpdate) {
   EXPECT_EQ(displays[1].id(), screen->GetPrimaryDisplay().id());
 }
 
+TEST_F(ScreenMusTest, SetDisplayForNewWindows) {
+  ScreenMus* screen = static_cast<ScreenMus*>(display::Screen::GetScreen());
+
+  // Set up 2 displays with display 1 as the display for new windows.
+  constexpr int64_t kDisplayId1 = 111;
+  constexpr int64_t kDisplayId2 = 222;
+  std::vector<display::Display> displays = {display::Display(kDisplayId1),
+                                            display::Display(kDisplayId2)};
+  screen->OnDisplaysChanged(ConvertDisplayToWsDisplays(displays), kDisplayId1,
+                            kDisplayId1,
+                            kDisplayId1 /* display_id_for_new_windows */);
+  EXPECT_EQ(kDisplayId1, screen->GetDisplayForNewWindows().id());
+
+  // Set display 2 as the display for new windows.
+  screen->OnDisplaysChanged(ConvertDisplayToWsDisplays(displays), kDisplayId1,
+                            kDisplayId1,
+                            kDisplayId2 /* display_id_for_new_windows */);
+  EXPECT_EQ(kDisplayId2, screen->GetDisplayForNewWindows().id());
+
+  // Set a bad display as the display for new windows. ScreenMus should fall
+  // back to the primary display.
+  screen->OnDisplaysChanged(ConvertDisplayToWsDisplays(displays), kDisplayId1,
+                            kDisplayId1, 666 /* display_id_for_new_windows */);
+  EXPECT_EQ(kDisplayId1, screen->GetDisplayForNewWindows().id());
+}
+
 }  // namespace
 }  // namespace views