ui-service: updates AuraTest* to more correctly do the right thing

Prior to this patch AuraTestHelper configured TestScreen in such a way
that it would create a WindowTreeHost to represent a screen. The
WindowTreeHost was created was not a WindowTreeHostMus, rather a
WindowTreeHostPlatform. This doesn't make sense for mus.

As part of this I'm also moving how ash configures Env to
AuraTestBase. This way other code using AuraTestBase can work better
with mus.

Lastly I needed to make sure sure AuraTestHelper resets some state
that was persisting between runs. This is necessitated by Env being
created at the TestSuite level, not per test.

BUG=776514
TEST=covered by tests

Change-Id: Icb1b5342f962f44c03185a3e40097491532c4efe
Reviewed-on: https://chromium-review.googlesource.com/791821Reviewed-by: Elliot Glaysher <erg@chromium.org>
Commit-Queue: Scott Violet <sky@chromium.org>
Cr-Commit-Position: refs/heads/master@{#519758}

ash/test/ash_test_suite.cc

@@ -4,9 +4,6 @@
 
 #include "ash/test/ash_test_suite.h"
 
-#include <memory>
-#include <set>
-
 #include "ash/public/cpp/config.h"
 #include "ash/test/ash_test_environment.h"
 #include "ash/test/ash_test_helper.h"
@@ -16,80 +13,14 @@
 #include "base/metrics/statistics_recorder.h"
 #include "base/path_service.h"
 #include "build/build_config.h"
-#include "cc/test/fake_output_surface.h"
-#include "cc/test/test_context_provider.h"
-#include "components/viz/test/test_layer_tree_frame_sink.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "ui/aura/test/aura_test_context_factory.h"
 #include "ui/base/resource/resource_bundle.h"
 #include "ui/base/ui_base_paths.h"
-#include "ui/compositor/test/fake_context_factory.h"
 #include "ui/gfx/gfx_paths.h"
 #include "ui/gl/test/gl_surface_test_support.h"
 
 namespace ash {
-namespace {
-
-class FrameSinkClient : public viz::TestLayerTreeFrameSinkClient {
- public:
-  explicit FrameSinkClient(
-      scoped_refptr<viz::ContextProvider> display_context_provider)
-      : display_context_provider_(std::move(display_context_provider)) {}
-
-  std::unique_ptr<viz::OutputSurface> CreateDisplayOutputSurface(
-      scoped_refptr<viz::ContextProvider> compositor_context_provider)
-      override {
-    return cc::FakeOutputSurface::Create3d(
-        std::move(display_context_provider_));
-  }
-
-  void DisplayReceivedLocalSurfaceId(
-      const viz::LocalSurfaceId& local_surface_id) override {}
-  void DisplayReceivedCompositorFrame(
-      const viz::CompositorFrame& frame) override {}
-  void DisplayWillDrawAndSwap(
-      bool will_draw_and_swap,
-      const viz::RenderPassList& render_passes) override {}
-  void DisplayDidDrawAndSwap() override {}
-
- private:
-  scoped_refptr<viz::ContextProvider> display_context_provider_;
-
-  DISALLOW_COPY_AND_ASSIGN(FrameSinkClient);
-};
-
-class AshTestContextFactory : public ui::FakeContextFactory {
- public:
-  AshTestContextFactory() = default;
-  ~AshTestContextFactory() override = default;
-
-  // ui::FakeContextFactory
-  void CreateLayerTreeFrameSink(
-      base::WeakPtr<ui::Compositor> compositor) override {
-    scoped_refptr<cc::TestContextProvider> context_provider =
-        cc::TestContextProvider::Create();
-    std::unique_ptr<FrameSinkClient> frame_sink_client =
-        std::make_unique<FrameSinkClient>(context_provider);
-    constexpr bool synchronous_composite = false;
-    constexpr bool disable_display_vsync = false;
-    const double refresh_rate = GetRefreshRate();
-    auto frame_sink = std::make_unique<viz::TestLayerTreeFrameSink>(
-        context_provider, cc::TestContextProvider::CreateWorker(), nullptr,
-        GetGpuMemoryBufferManager(), renderer_settings(),
-        base::ThreadTaskRunnerHandle::Get().get(), synchronous_composite,
-        disable_display_vsync, refresh_rate);
-    frame_sink->SetClient(frame_sink_client.get());
-    compositor->SetLayerTreeFrameSink(std::move(frame_sink));
-    frame_sink_clients_.insert(std::move(frame_sink_client));
-  }
-
- private:
-  std::set<std::unique_ptr<viz::TestLayerTreeFrameSinkClient>>
-      frame_sink_clients_;
-
-  DISALLOW_COPY_AND_ASSIGN(AshTestContextFactory);
-};
-
-}  // namespace
 
 AshTestSuite::AshTestSuite(int argc, char** argv) : TestSuite(argc, argv) {}
 
@@ -139,7 +70,7 @@ void AshTestSuite::Initialize() {
                                                      : aura::Env::Mode::LOCAL);
 
   if (is_mus || is_mash) {
-    context_factory_ = std::make_unique<AshTestContextFactory>();
+    context_factory_ = std::make_unique<aura::test::AuraTestContextFactory>();
     env_->set_context_factory(context_factory_.get());
     env_->set_context_factory_private(nullptr);
   }

components/exo/BUILD.gn

@@ -190,6 +190,7 @@ test("exo_unittests") {
     "//mojo/edk/embedder:headers",
     "//testing/gtest",
     "//ui/aura",
+    "//ui/aura:test_support",
     "//ui/base",
     "//ui/compositor:test_support",
     "//ui/gl:test_support",

components/exo/wayland/BUILD.gn

@@ -322,6 +322,7 @@ test("wayland_client_perftests") {
     "//mojo/edk/embedder:headers",
     "//testing/perf",
     "//ui/aura",
+    "//ui/aura:test_support",
     "//ui/base",
     "//ui/compositor:test_support",
     "//ui/events:gesture_detection",

ui/aura/BUILD.gn

@@ -223,6 +223,8 @@ static_library("test_support") {
     "test/aura_mus_test_base.h",
     "test/aura_test_base.cc",
     "test/aura_test_base.h",
+    "test/aura_test_context_factory.cc",
+    "test/aura_test_context_factory.h",
     "test/aura_test_helper.cc",
     "test/aura_test_helper.h",
     "test/aura_test_utils.cc",
@@ -270,6 +272,8 @@ static_library("test_support") {
   ]
   deps = [
     "//base/test:test_support",
+    "//cc:test_support",
+    "//components/viz/test:test_support",
     "//services/ui/public/interfaces",
     "//skia",
     "//testing/gtest",

ui/aura/mus/window_tree_client.cc

@@ -208,7 +208,9 @@ WindowTreeClient::WindowTreeClient(
   // Allow for a null request in tests.
   if (request.is_pending())
     binding_.Bind(std::move(request));
-  client::GetTransientWindowClient()->AddObserver(this);
+  // Some tests may not create a TransientWindowClient.
+  if (client::GetTransientWindowClient())
+    client::GetTransientWindowClient()->AddObserver(this);
   if (window_manager_delegate)
     window_manager_delegate->SetWindowManagerClient(this);
   if (connector) {  // |connector| can be null in tests.
@@ -253,7 +255,9 @@ WindowTreeClient::~WindowTreeClient() {
 
   capture_synchronizer_.reset();
 
-  client::GetTransientWindowClient()->RemoveObserver(this);
+  // Some tests may not create a TransientWindowClient.
+  if (client::GetTransientWindowClient())
+    client::GetTransientWindowClient()->RemoveObserver(this);
 
   Env* env = Env::GetInstance();
   if (compositor_context_factory_ &&
@@ -515,9 +519,13 @@ void WindowTreeClient::CreateOrUpdateWindowFromWindowData(
   if (window_data.transient_parent_id == kInvalidServerId)
     return;
 
-  // Adjust the transient parent if necessary.
+  // Some tests may not create a TransientWindowClient.
   client::TransientWindowClient* transient_window_client =
       client::GetTransientWindowClient();
+  if (!transient_window_client)
+    return;
+
+  // Adjust the transient parent if necessary.
   Window* existing_transient_parent =
       transient_window_client->GetTransientParent(window->GetWindow());
   WindowMus* new_transient_parent =

ui/aura/mus/window_tree_client_unittest.cc

@@ -67,7 +67,7 @@ const char kTestPropertyServerKey2[] = "test-property-server2";
 const char kTestPropertyServerKey3[] = "test-property-server3";
 
 Id server_id(Window* window) {
-  return WindowMus::Get(window)->server_id();
+  return window ? WindowMus::Get(window)->server_id() : 0;
 }
 
 std::unique_ptr<Window> CreateWindowUsingId(
@@ -137,7 +137,51 @@ WindowTreeHostMusInitParams CreateWindowTreeHostMusInitParams(
 }  // namespace
 
 using WindowTreeClientWmTest = test::AuraMusWmTestBase;
-using WindowTreeClientClientTest = test::AuraMusClientTestBase;
+
+class WindowTreeClientClientTest : public test::AuraMusClientTestBase {
+ public:
+  WindowTreeClientClientTest() = default;
+  ~WindowTreeClientClientTest() override = default;
+
+  struct TopLevel {
+    std::unique_ptr<client::DefaultCaptureClient> capture_client;
+    std::unique_ptr<WindowTreeHostMus> host;
+  };
+
+  std::unique_ptr<TopLevel> CreateWindowTreeHostForTopLevel() {
+    std::unique_ptr<TopLevel> top_level = std::make_unique<TopLevel>();
+    top_level->host = std::make_unique<WindowTreeHostMus>(
+        CreateInitParamsForTopLevel(window_tree_client_impl()));
+    top_level->host->InitHost();
+    Window* top_level_window = top_level->host->window();
+    top_level->capture_client =
+        std::make_unique<client::DefaultCaptureClient>();
+    client::SetCaptureClient(top_level_window, top_level->capture_client.get());
+    window_tree_client_impl()->capture_synchronizer()->AttachToCaptureClient(
+        top_level->capture_client.get());
+
+    // Ack the request to the windowtree to create the new window.
+    uint32_t change_id = 0;
+    EXPECT_TRUE(window_tree()->GetAndRemoveFirstChangeOfType(
+        WindowTreeChangeType::NEW_TOP_LEVEL, &change_id));
+    EXPECT_EQ(window_tree()->window_id(), server_id(top_level_window));
+
+    ui::mojom::WindowDataPtr data = ui::mojom::WindowData::New();
+    data->window_id = server_id(top_level_window);
+    data->visible = true;
+    window_tree_client()->OnTopLevelCreated(
+        change_id, std::move(data), next_display_id_++, true, base::nullopt);
+    EXPECT_EQ(0u, window_tree()->GetChangeCountForType(
+                      WindowTreeChangeType::VISIBLE));
+    EXPECT_TRUE(top_level_window->TargetVisibility());
+    return top_level;
+  }
+
+ private:
+  int64_t next_display_id_ = 1;
+
+  DISALLOW_COPY_AND_ASSIGN(WindowTreeClientClientTest);
+};
 
 class WindowTreeClientWmTestSurfaceSync
     : public WindowTreeClientWmTest,
@@ -546,65 +590,69 @@ TEST_F(WindowTreeClientWmTest, FocusFromServer) {
 // server replies with a new bounds and the original bounds change fails.
 // The server bounds change takes hold along with the associated
 // viz::LocalSurfaceId.
-TEST_F(WindowTreeClientWmTest, SetBoundsFailedWithPendingChange) {
-  const gfx::Rect original_bounds(root_window()->bounds());
+TEST_F(WindowTreeClientClientTest, SetBoundsFailedWithPendingChange) {
+  aura::Window root_window(nullptr);
+  root_window.Init(ui::LAYER_NOT_DRAWN);
+  const gfx::Rect original_bounds(root_window.bounds());
   const gfx::Rect new_bounds(gfx::Rect(0, 0, 100, 100));
-  ASSERT_NE(new_bounds, root_window()->bounds());
-  root_window()->SetBounds(new_bounds);
-  EXPECT_EQ(new_bounds, root_window()->bounds());
+  ASSERT_NE(new_bounds, root_window.bounds());
+  root_window.SetBounds(new_bounds);
+  EXPECT_EQ(new_bounds, root_window.bounds());
 
   // Simulate the server responding with a bounds change.
   const gfx::Rect server_changed_bounds(gfx::Rect(0, 0, 101, 102));
   const viz::LocalSurfaceId server_changed_local_surface_id(
       1, base::UnguessableToken::Create());
   window_tree_client()->OnWindowBoundsChanged(
-      server_id(root_window()), original_bounds, server_changed_bounds,
+      server_id(&root_window), original_bounds, server_changed_bounds,
       server_changed_local_surface_id);
 
-  WindowMus* root_window_mus = WindowMus::Get(root_window());
+  WindowMus* root_window_mus = WindowMus::Get(&root_window);
   ASSERT_NE(nullptr, root_window_mus);
 
   // This shouldn't trigger the bounds changing yet.
-  EXPECT_EQ(new_bounds, root_window()->bounds());
+  EXPECT_EQ(new_bounds, root_window.bounds());
   EXPECT_FALSE(root_window_mus->GetLocalSurfaceId().is_valid());
 
   // Tell the client the change failed, which should trigger failing to the
   // most recent bounds from server.
   ASSERT_TRUE(window_tree()->AckSingleChangeOfType(WindowTreeChangeType::BOUNDS,
                                                    false));
-  EXPECT_EQ(server_changed_bounds, root_window()->bounds());
+  EXPECT_EQ(server_changed_bounds, root_window.bounds());
   EXPECT_EQ(server_changed_local_surface_id,
             root_window_mus->GetLocalSurfaceId());
 
   // Simulate server changing back to original bounds. Should take immediately.
-  window_tree_client()->OnWindowBoundsChanged(server_id(root_window()),
+  window_tree_client()->OnWindowBoundsChanged(server_id(&root_window),
                                               server_changed_bounds,
                                               original_bounds, base::nullopt);
-  EXPECT_EQ(original_bounds, root_window()->bounds());
+  EXPECT_EQ(original_bounds, root_window.bounds());
   EXPECT_FALSE(root_window_mus->GetLocalSurfaceId().is_valid());
 }
 
-TEST_F(WindowTreeClientWmTest, TwoInFlightBoundsChangesBothCanceled) {
-  const gfx::Rect original_bounds(root_window()->bounds());
+TEST_F(WindowTreeClientClientTest, TwoInFlightBoundsChangesBothCanceled) {
+  aura::Window root_window(nullptr);
+  root_window.Init(ui::LAYER_NOT_DRAWN);
+  const gfx::Rect original_bounds(root_window.bounds());
   const gfx::Rect bounds1(gfx::Rect(0, 0, 100, 100));
   const gfx::Rect bounds2(gfx::Rect(0, 0, 100, 102));
-  root_window()->SetBounds(bounds1);
-  EXPECT_EQ(bounds1, root_window()->bounds());
+  root_window.SetBounds(bounds1);
+  EXPECT_EQ(bounds1, root_window.bounds());
 
-  root_window()->SetBounds(bounds2);
-  EXPECT_EQ(bounds2, root_window()->bounds());
+  root_window.SetBounds(bounds2);
+  EXPECT_EQ(bounds2, root_window.bounds());
 
   // Tell the client the first bounds failed. As there is a still a change in
   // flight nothing should happen.
   ASSERT_TRUE(
       window_tree()->AckFirstChangeOfType(WindowTreeChangeType::BOUNDS, false));
-  EXPECT_EQ(bounds2, root_window()->bounds());
+  EXPECT_EQ(bounds2, root_window.bounds());
 
   // Tell the client the seconds bounds failed. Should now fallback to original
   // value.
   ASSERT_TRUE(window_tree()->AckSingleChangeOfType(WindowTreeChangeType::BOUNDS,
                                                    false));
-  EXPECT_EQ(original_bounds, root_window()->bounds());
+  EXPECT_EQ(original_bounds, root_window.bounds());
 }
 
 TEST_F(WindowTreeClientWmTest, TwoInFlightTransformsChangesBothCanceled) {
@@ -1646,7 +1694,6 @@ TEST_F(WindowTreeClientClientTest, NewTopLevelWindow) {
   window_tree_host->InitHost();
   EXPECT_FALSE(window_tree_host->window()->TargetVisibility());
   aura::Window* top_level = window_tree_host->window();
-  EXPECT_NE(server_id(top_level), server_id(root_window()));
   EXPECT_EQ(initial_root_count + 1,
             window_tree_client_impl()->GetRoots().size());
   EXPECT_TRUE(window_tree_client_impl()->GetRoots().count(top_level) > 0u);
@@ -1809,14 +1856,16 @@ TEST_F(WindowTreeClientClientTest, NewWindowGetsProperties) {
 
 // Assertions around transient windows.
 TEST_F(WindowTreeClientClientTest, Transients) {
+  aura::Window root_window(nullptr);
+  root_window.Init(ui::LAYER_NOT_DRAWN);
   client::TransientWindowClient* transient_client =
       client::GetTransientWindowClient();
   Window parent(nullptr);
   parent.Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(&parent);
+  root_window.AddChild(&parent);
   Window transient(nullptr);
   transient.Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(&transient);
+  root_window.AddChild(&transient);
   window_tree()->AckAllChanges();
   transient_client->AddTransientChild(&parent, &transient);
   ASSERT_EQ(1u, window_tree()->GetChangeCountForType(
@@ -1885,23 +1934,25 @@ TEST_F(WindowTreeClientClientTest, DontRestackTransientsFromOtherClients) {
 // when the change originates from the server.
 TEST_F(WindowTreeClientClientTest,
        TransientChildServerMutateNotifiesOfRestack) {
+  aura::Window root_window(nullptr);
+  root_window.Init(ui::LAYER_NOT_DRAWN);
   Window* w1 = new Window(nullptr);
   w1->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w1);
+  root_window.AddChild(w1);
   Window* w2 = new Window(nullptr);
   w2->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w2);
+  root_window.AddChild(w2);
   Window* w3 = new Window(nullptr);
   w3->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w3);
+  root_window.AddChild(w3);
   // Three children of root: |w1|, |w2| and |w3| (in that order). Make |w1| a
   // transient child of |w2|. Should trigger moving |w1| on top of |w2|, but not
   // notify the server of the reorder.
   window_tree()->AckAllChanges();
   window_tree_client()->OnTransientWindowAdded(server_id(w2), server_id(w1));
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w1, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w1, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   // Only reorders should be generated.
   EXPECT_NE(0u, window_tree()->number_of_changes());
   window_tree()->AckAllChangesOfType(WindowTreeChangeType::REORDER, true);
@@ -1909,9 +1960,9 @@ TEST_F(WindowTreeClientClientTest,
 
   // Make |w3| also a transient child of |w2|.
   window_tree_client()->OnTransientWindowAdded(server_id(w2), server_id(w3));
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w1, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w1, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   // Only reorders should be generated.
   EXPECT_NE(0u, window_tree()->number_of_changes());
   window_tree()->AckAllChangesOfType(WindowTreeChangeType::REORDER, true);
@@ -1919,9 +1970,9 @@ TEST_F(WindowTreeClientClientTest,
 
   // Remove |w1| as a transient child, this should move |w3| on top of |w2|.
   window_tree_client()->OnTransientWindowRemoved(server_id(w2), server_id(w1));
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w3, root_window()->children()[1]);
-  EXPECT_EQ(w1, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w3, root_window.children()[1]);
+  EXPECT_EQ(w1, root_window.children()[2]);
   // Only reorders should be generated.
   EXPECT_NE(0u, window_tree()->number_of_changes());
   window_tree()->AckAllChangesOfType(WindowTreeChangeType::REORDER, true);
@@ -1932,25 +1983,28 @@ TEST_F(WindowTreeClientClientTest,
 // restacks;
 TEST_F(WindowTreeClientClientTest,
        TransientChildClientMutateNotifiesOfRestack) {
+  aura::Window root_window(nullptr);
+  root_window.Init(ui::LAYER_NOT_DRAWN);
+
   client::TransientWindowClient* transient_client =
       client::GetTransientWindowClient();
   Window* w1 = new Window(nullptr);
   w1->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w1);
+  root_window.AddChild(w1);
   Window* w2 = new Window(nullptr);
   w2->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w2);
+  root_window.AddChild(w2);
   Window* w3 = new Window(nullptr);
   w3->Init(ui::LAYER_NOT_DRAWN);
-  root_window()->AddChild(w3);
+  root_window.AddChild(w3);
   // Three children of root: |w1|, |w2| and |w3| (in that order). Make |w1| a
   // transient child of |w2|. Should trigger moving |w1| on top of |w2|, and
   // notify notify the server of the reorder.
   window_tree()->AckAllChanges();
   transient_client->AddTransientChild(w2, w1);
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w1, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w1, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
       WindowTreeChangeType::ADD_TRANSIENT, true));
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
@@ -1959,9 +2013,9 @@ TEST_F(WindowTreeClientClientTest,
 
   // Make |w3| also a transient child of |w2|. Order shouldn't change.
   transient_client->AddTransientChild(w2, w3);
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w1, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w1, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
       WindowTreeChangeType::ADD_TRANSIENT, true));
   // While the order doesn't change, internally aura shuffles things around,
@@ -1972,9 +2026,9 @@ TEST_F(WindowTreeClientClientTest,
 
   // Remove |w1| as a transient child, this should move |w3| on top of |w2|.
   transient_client->RemoveTransientChild(w2, w1);
-  EXPECT_EQ(w2, root_window()->children()[0]);
-  EXPECT_EQ(w3, root_window()->children()[1]);
-  EXPECT_EQ(w1, root_window()->children()[2]);
+  EXPECT_EQ(w2, root_window.children()[0]);
+  EXPECT_EQ(w3, root_window.children()[1]);
+  EXPECT_EQ(w1, root_window.children()[2]);
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
       WindowTreeChangeType::REMOVE_TRANSIENT, true));
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
@@ -1982,20 +2036,20 @@ TEST_F(WindowTreeClientClientTest,
   EXPECT_EQ(0u, window_tree()->number_of_changes());
 
   // Make |w1| the first child and ensure a REORDER was scheduled.
-  root_window()->StackChildAtBottom(w1);
-  EXPECT_EQ(w1, root_window()->children()[0]);
-  EXPECT_EQ(w2, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  root_window.StackChildAtBottom(w1);
+  EXPECT_EQ(w1, root_window.children()[0]);
+  EXPECT_EQ(w2, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   EXPECT_TRUE(window_tree()->AckSingleChangeOfType(
       WindowTreeChangeType::REORDER, true));
   EXPECT_EQ(0u, window_tree()->number_of_changes());
 
   // Try stacking |w2| above |w3|. This should be disallowed as that would
   // result in placing |w2| above its transient child.
-  root_window()->StackChildAbove(w2, w3);
-  EXPECT_EQ(w1, root_window()->children()[0]);
-  EXPECT_EQ(w2, root_window()->children()[1]);
-  EXPECT_EQ(w3, root_window()->children()[2]);
+  root_window.StackChildAbove(w2, w3);
+  EXPECT_EQ(w1, root_window.children()[0]);
+  EXPECT_EQ(w2, root_window.children()[1]);
+  EXPECT_EQ(w3, root_window.children()[2]);
   // The stack above is followed by a reorder from TransientWindowManager,
   // hence multiple changes.
   EXPECT_NE(0u, window_tree()->number_of_changes());
@@ -2310,76 +2364,50 @@ TEST_F(WindowTreeClientWmTest, OnWindowTreeCaptureChanged) {
 }
 
 TEST_F(WindowTreeClientClientTest, TwoWindowTreesRequestCapture) {
-  // Creating a WindowTreeHost so we can have two root windows: top_level
-  // and root_window().
-  std::unique_ptr<WindowTreeHostMus> window_tree_host =
-      std::make_unique<WindowTreeHostMus>(
-          CreateInitParamsForTopLevel(window_tree_client_impl()));
-  window_tree_host->InitHost();
-  Window* top_level = window_tree_host->window();
-  std::unique_ptr<client::DefaultCaptureClient> capture_client(
-      std::make_unique<client::DefaultCaptureClient>());
-  client::SetCaptureClient(top_level, capture_client.get());
-  window_tree_client_impl()->capture_synchronizer()->AttachToCaptureClient(
-      capture_client.get());
-  EXPECT_NE(server_id(top_level), server_id(root_window()));
-
-  // Ack the request to the windowtree to create the new window.
-  uint32_t change_id;
-  ASSERT_TRUE(window_tree()->GetAndRemoveFirstChangeOfType(
-      WindowTreeChangeType::NEW_TOP_LEVEL, &change_id));
-  EXPECT_EQ(window_tree()->window_id(), server_id(top_level));
-
-  ui::mojom::WindowDataPtr data = ui::mojom::WindowData::New();
-  data->window_id = server_id(top_level);
-  data->visible = true;
-  const int64_t display_id = 1;
-  window_tree_client()->OnTopLevelCreated(change_id, std::move(data),
-                                          display_id, true, base::nullopt);
-  EXPECT_EQ(
-      0u, window_tree()->GetChangeCountForType(WindowTreeChangeType::VISIBLE));
-  EXPECT_TRUE(top_level->TargetVisibility());
+  std::unique_ptr<TopLevel> top_level1 = CreateWindowTreeHostForTopLevel();
+  std::unique_ptr<TopLevel> top_level2 = CreateWindowTreeHostForTopLevel();
 
+  aura::Window* root_window1 = top_level1->host->window();
+  aura::Window* root_window2 = top_level2->host->window();
   std::unique_ptr<CaptureRecorder> capture_recorder1(
-      std::make_unique<CaptureRecorder>(root_window()));
+      std::make_unique<CaptureRecorder>(root_window1));
   std::unique_ptr<CaptureRecorder> capture_recorder2(
-      std::make_unique<CaptureRecorder>(top_level));
-  EXPECT_NE(client::GetCaptureClient(root_window()),
-            client::GetCaptureClient(top_level));
+      std::make_unique<CaptureRecorder>(root_window2));
+  EXPECT_NE(client::GetCaptureClient(root_window1),
+            client::GetCaptureClient(root_window2));
 
   EXPECT_EQ(0, capture_recorder1->capture_changed_count());
   EXPECT_EQ(0, capture_recorder2->capture_changed_count());
-  // Give capture to top_level and ensure everyone is notified correctly.
-  top_level->SetCapture();
+  // Give capture to root_window2 and ensure everyone is notified correctly.
+  root_window2->SetCapture();
   ASSERT_TRUE(window_tree()->AckSingleChangeOfType(
       WindowTreeChangeType::CAPTURE, true));
   EXPECT_EQ(0, capture_recorder1->capture_changed_count());
   EXPECT_EQ(1, capture_recorder2->capture_changed_count());
-  EXPECT_EQ(top_level->id(),
+  EXPECT_EQ(root_window2->id(),
             capture_recorder2->last_gained_capture_window_id());
   EXPECT_EQ(0, capture_recorder2->last_lost_capture_window_id());
-  top_level->ReleaseCapture();
+  root_window2->ReleaseCapture();
   capture_recorder1->reset_capture_captured_count();
   capture_recorder2->reset_capture_captured_count();
 
-  // Release capture of top_level shouldn't affect the capture of root_window().
-  top_level->SetCapture();
-  root_window()->SetCapture();
-  top_level->ReleaseCapture();
+  // Release capture of  shouldn't affect the capture of root_window1.
+  root_window2->SetCapture();
+  root_window1->SetCapture();
+  root_window2->ReleaseCapture();
   EXPECT_EQ(1, capture_recorder1->capture_changed_count());
   EXPECT_EQ(2, capture_recorder2->capture_changed_count());
-  EXPECT_EQ(root_window()->id(),
+  EXPECT_EQ(root_window1->id(),
             capture_recorder1->last_gained_capture_window_id());
   EXPECT_EQ(0, capture_recorder1->last_lost_capture_window_id());
   EXPECT_EQ(0, capture_recorder2->last_gained_capture_window_id());
-  EXPECT_EQ(top_level->id(), capture_recorder2->last_lost_capture_window_id());
+  EXPECT_EQ(root_window2->id(),
+            capture_recorder2->last_lost_capture_window_id());
 
   capture_recorder1->reset_capture_captured_count();
   capture_recorder2->reset_capture_captured_count();
   capture_recorder1.reset();
   capture_recorder2.reset();
-  window_tree_host.reset();
-  capture_client.reset();
 }
 
 TEST_F(WindowTreeClientClientTest, ModalTypeWindowFail) {

ui/aura/test/DEPS

 include_rules = [
+  "+cc/test",
+  "+components/viz/test",
   "+mojo/edk/embedder/embedder.h",
   "+ui/gl",
   "+ui/wm/core/wm_state.h",

ui/aura/test/aura_mus_test_base.cc

@@ -16,16 +16,6 @@ void AuraMusWmTestBase::SetUp() {
   // in env.
   EnableMusWithTestWindowTree();
   AuraTestBase::SetUp();
-
-  aura::Env* env = aura::Env::GetInstance();
-  DCHECK(env);
-  context_factory_to_restore_ = env->context_factory();
-  env->set_context_factory(&context_factory_);
-}
-
-void AuraMusWmTestBase::TearDown() {
-  aura::Env::GetInstance()->set_context_factory(context_factory_to_restore_);
-  AuraTestBase::TearDown();
 }
 
 AuraMusClientTestBase::AuraMusClientTestBase() {}
@@ -38,16 +28,6 @@ void AuraMusClientTestBase::SetUp() {
   EnableMusWithTestWindowTree();
   set_window_manager_delegate(nullptr);
   AuraTestBase::SetUp();
-
-  aura::Env* env = aura::Env::GetInstance();
-  DCHECK(env);
-  context_factory_to_restore_ = env->context_factory();
-  env->set_context_factory(&context_factory_);
-}
-
-void AuraMusClientTestBase::TearDown() {
-  aura::Env::GetInstance()->set_context_factory(context_factory_to_restore_);
-  AuraTestBase::TearDown();
 }
 
 }  // namespace test

ui/aura/test/aura_mus_test_base.h

@@ -7,7 +7,6 @@
 
 #include "base/macros.h"
 #include "ui/aura/test/aura_test_base.h"
-#include "ui/compositor/test/fake_context_factory.h"
 
 namespace aura {
 namespace test {
@@ -24,12 +23,8 @@ class AuraMusWmTestBase : public AuraTestBase {
 
   // AuraTestBase:
   void SetUp() override;
-  void TearDown() override;
 
  private:
-  ui::FakeContextFactory context_factory_;
-  ui::ContextFactory* context_factory_to_restore_ = nullptr;
-
   DISALLOW_COPY_AND_ASSIGN(AuraMusWmTestBase);
 };
 
@@ -46,12 +41,8 @@ class AuraMusClientTestBase : public AuraTestBase {
 
   // AuraTestBase:
   void SetUp() override;
-  void TearDown() override;
 
  private:
-  ui::FakeContextFactory context_factory_;
-  ui::ContextFactory* context_factory_to_restore_ = nullptr;
-
   DISALLOW_COPY_AND_ASSIGN(AuraMusClientTestBase);
 };
 

ui/aura/test/aura_test_base.cc

@@ -9,6 +9,7 @@
 #include "ui/aura/mus/property_utils.h"
 #include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/mus/window_tree_host_mus.h"
+#include "ui/aura/test/aura_test_context_factory.h"
 #include "ui/aura/test/test_window_delegate.h"
 #include "ui/aura/window.h"
 #include "ui/base/ime/input_method_initializer.h"
@@ -77,11 +78,16 @@ void AuraTestBase::SetUp() {
       ui::VelocityTracker::Strategy::LSQ2_RESTRICTED);
 
   // The ContextFactory must exist before any Compositors are created.
-  bool enable_pixel_output = false;
   ui::ContextFactory* context_factory = nullptr;
   ui::ContextFactoryPrivate* context_factory_private = nullptr;
-  ui::InitializeContextFactoryForTests(enable_pixel_output, &context_factory,
-                                       &context_factory_private);
+  if (use_mus_) {
+    mus_context_factory_ = std::make_unique<AuraTestContextFactory>();
+    context_factory = mus_context_factory_.get();
+  } else {
+    const bool enable_pixel_output = false;
+    ui::InitializeContextFactoryForTests(enable_pixel_output, &context_factory,
+                                         &context_factory_private);
+  }
 
   helper_ = std::make_unique<AuraTestHelper>();
   if (use_mus_) {
@@ -98,9 +104,18 @@ void AuraTestBase::TearDown() {
   // and these tasks if un-executed would upset Valgrind.
   RunAllPendingInMessageLoop();
 
-  window_tree_hosts_.clear();
+  // AuraTestHelper may own a WindowTreeHost, don't delete it here else
+  // AuraTestHelper will have use after frees.
+  for (size_t i = window_tree_hosts_.size(); i > 0; --i) {
+    if (window_tree_hosts_[i - 1].get() == helper_->host()) {
+      window_tree_hosts_[i - 1].release();
+      window_tree_hosts_.erase(window_tree_hosts_.begin() + i - 1);
+      break;
+    }
+  }
 
   helper_->TearDown();
+  window_tree_hosts_.clear();
   ui::TerminateContextFactoryForTests();
   ui::ShutdownInputMethodForTesting();
   testing::Test::TearDown();

ui/aura/test/aura_test_base.h

@@ -35,6 +35,8 @@ class FocusClient;
 
 namespace test {
 
+class AuraTestContextFactory;
+
 enum class BackendType { CLASSIC, MUS };
 
 // A base class for aura unit tests.
@@ -169,6 +171,7 @@ class AuraTestBase : public testing::Test,
   bool teardown_called_ = false;
   PropertyConverter property_converter_;
   std::unique_ptr<AuraTestHelper> helper_;
+  std::unique_ptr<AuraTestContextFactory> mus_context_factory_;
   std::vector<std::unique_ptr<WindowTreeHostMus>> window_tree_hosts_;
   std::vector<std::unique_ptr<ui::PointerEvent>> observed_pointer_events_;
 

ui/aura/test/aura_test_context_factory.cc

+// Copyright 2017 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 "ui/aura/test/aura_test_context_factory.h"
+
+#include "cc/test/fake_output_surface.h"
+#include "cc/test/test_context_provider.h"
+#include "components/viz/test/test_layer_tree_frame_sink.h"
+
+namespace aura {
+namespace test {
+namespace {
+
+class FrameSinkClient : public viz::TestLayerTreeFrameSinkClient {
+ public:
+  explicit FrameSinkClient(
+      scoped_refptr<viz::ContextProvider> display_context_provider)
+      : display_context_provider_(std::move(display_context_provider)) {}
+
+  // viz::TestLayerTreeFrameSinkClient:
+  std::unique_ptr<viz::OutputSurface> CreateDisplayOutputSurface(
+      scoped_refptr<viz::ContextProvider> compositor_context_provider)
+      override {
+    return cc::FakeOutputSurface::Create3d(
+        std::move(display_context_provider_));
+  }
+  void DisplayReceivedLocalSurfaceId(
+      const viz::LocalSurfaceId& local_surface_id) override {}
+  void DisplayReceivedCompositorFrame(
+      const viz::CompositorFrame& frame) override {}
+  void DisplayWillDrawAndSwap(
+      bool will_draw_and_swap,
+      const viz::RenderPassList& render_passes) override {}
+  void DisplayDidDrawAndSwap() override {}
+
+ private:
+  scoped_refptr<viz::ContextProvider> display_context_provider_;
+
+  DISALLOW_COPY_AND_ASSIGN(FrameSinkClient);
+};
+
+}  // namespace
+
+AuraTestContextFactory::AuraTestContextFactory() = default;
+
+AuraTestContextFactory::~AuraTestContextFactory() = default;
+
+void AuraTestContextFactory::CreateLayerTreeFrameSink(
+    base::WeakPtr<ui::Compositor> compositor) {
+  scoped_refptr<cc::TestContextProvider> context_provider =
+      cc::TestContextProvider::Create();
+  std::unique_ptr<FrameSinkClient> frame_sink_client =
+      std::make_unique<FrameSinkClient>(context_provider);
+  constexpr bool synchronous_composite = false;
+  constexpr bool disable_display_vsync = false;
+  const double refresh_rate = GetRefreshRate();
+  auto frame_sink = std::make_unique<viz::TestLayerTreeFrameSink>(
+      context_provider, cc::TestContextProvider::CreateWorker(), nullptr,
+      GetGpuMemoryBufferManager(), renderer_settings(),
+      base::ThreadTaskRunnerHandle::Get().get(), synchronous_composite,
+      disable_display_vsync, refresh_rate);
+  frame_sink->SetClient(frame_sink_client.get());
+  compositor->SetLayerTreeFrameSink(std::move(frame_sink));
+  frame_sink_clients_.insert(std::move(frame_sink_client));
+}
+
+}  // namespace test
+}  // namespace aura

ui/aura/test/aura_test_context_factory.h

+// Copyright 2017 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 UI_AURA_TEST_AURA_TEST_CONTEXT_FACTORY_H_
+#define UI_AURA_TEST_AURA_TEST_CONTEXT_FACTORY_H_
+
+#include <memory>
+
+#include "base/macros.h"
+#include "ui/compositor/test/fake_context_factory.h"
+
+namespace viz {
+class TestLayerTreeFrameSinkClient;
+}
+
+namespace aura {
+namespace test {
+
+class AuraTestContextFactory : public ui::FakeContextFactory {
+ public:
+  AuraTestContextFactory();
+  ~AuraTestContextFactory() override;
+
+  // ui::FakeContextFactory
+  void CreateLayerTreeFrameSink(
+      base::WeakPtr<ui::Compositor> compositor) override;
+
+ private:
+  std::set<std::unique_ptr<viz::TestLayerTreeFrameSinkClient>>
+      frame_sink_clients_;
+
+  DISALLOW_COPY_AND_ASSIGN(AuraTestContextFactory);
+};
+
+}  // namespace test
+}  // namespace aura
+
+#endif  // UI_AURA_TEST_AURA_TEST_CONTEXT_FACTORY_H_<

ui/aura/test/aura_test_helper.cc

@@ -109,6 +109,8 @@ void AuraTestHelper::SetUp(ui::ContextFactory* context_factory,
     InitWindowTreeClient();
   if (!Env::GetInstanceDontCreate())
     env_ = Env::CreateInstance(env_mode);
+  else
+    env_mode_to_restore_ = Env::GetInstance()->mode();
   EnvTestHelper env_helper;
   // Always reset the mode. This really only matters for if Env was created
   // above.
@@ -117,6 +119,9 @@ void AuraTestHelper::SetUp(ui::ContextFactory* context_factory,
   // Tests assume they can set the mouse location on Env() and have it reflected
   // in tests.
   env_helper.SetAlwaysUseLastMouseLocation(true);
+  context_factory_to_restore_ = Env::GetInstance()->context_factory();
+  context_factory_private_to_restore_ =
+      Env::GetInstance()->context_factory_private();
   Env::GetInstance()->set_context_factory(context_factory);
   Env::GetInstance()->set_context_factory_private(context_factory_private);
   // Unit tests generally don't want to query the system, rather use the state
@@ -133,20 +138,22 @@ void AuraTestHelper::SetUp(ui::ContextFactory* context_factory,
     // This must be reset before creating TestScreen, which sets up the display
     // scale factor for this test iteration.
     display::Display::ResetForceDeviceScaleFactorForTesting();
-    test_screen_.reset(TestScreen::Create(host_size));
+    test_screen_.reset(TestScreen::Create(host_size, window_tree_client_));
     if (!screen)
       display::Screen::SetScreenInstance(test_screen_.get());
-    host_.reset(test_screen_->CreateHostForPrimaryDisplay());
-    host_->window()->SetEventTargeter(
-        std::unique_ptr<ui::EventTargeter>(new WindowTargeter()));
-
-    client::SetFocusClient(root_window(), focus_client_.get());
-    client::SetCaptureClient(root_window(), capture_client());
-    parenting_client_.reset(new TestWindowParentingClient(root_window()));
-
-    root_window()->Show();
-    // Ensure width != height so tests won't confuse them.
-    host()->SetBoundsInPixels(gfx::Rect(host_size));
+    if (env_mode == Env::Mode::LOCAL || window_manager_delegate_) {
+      host_.reset(test_screen_->CreateHostForPrimaryDisplay());
+      host_->window()->SetEventTargeter(
+          std::unique_ptr<ui::EventTargeter>(new WindowTargeter()));
+
+      client::SetFocusClient(root_window(), focus_client_.get());
+      client::SetCaptureClient(root_window(), capture_client());
+      parenting_client_.reset(new TestWindowParentingClient(root_window()));
+
+      root_window()->Show();
+      // Ensure width != height so tests won't confuse them.
+      host()->SetBoundsInPixels(gfx::Rect(host_size));
+    }
   }
 
   if (mode_ == Mode::MUS_CREATE_WINDOW_TREE_CLIENT) {
@@ -162,7 +169,7 @@ void AuraTestHelper::TearDown() {
   g_instance = nullptr;
   teardown_called_ = true;
   parenting_client_.reset();
-  if (mode_ != Mode::MUS) {
+  if (mode_ != Mode::MUS && root_window()) {
     client::SetFocusClient(root_window(), nullptr);
     client::SetCaptureClient(root_window(), nullptr);
     host_.reset();
@@ -174,12 +181,23 @@ void AuraTestHelper::TearDown() {
     window_tree_client_setup_.reset();
     focus_client_.reset();
     capture_client_.reset();
+  } else {
+    if (display::Screen::GetScreen() == test_screen_.get())
+      display::Screen::SetScreenInstance(nullptr);
+    test_screen_.reset();
+    window_tree_client_setup_.reset();
   }
   ui::GestureRecognizer::Reset();
   ui::ShutdownInputMethodForTesting();
 
-  if (env_)
+  if (env_) {
     env_.reset();
+  } else {
+    Env::GetInstance()->set_context_factory(context_factory_to_restore_);
+    Env::GetInstance()->set_context_factory_private(
+        context_factory_private_to_restore_);
+    EnvTestHelper().SetMode(env_mode_to_restore_);
+  }
   wm_state_.reset();
 }
 

ui/aura/test/aura_test_helper.h

@@ -14,6 +14,7 @@
 
 namespace ui {
 class ContextFactory;
+class ContextFactoryPrivate;
 class ScopedAnimationDurationScaleMode;
 }
 
@@ -76,7 +77,7 @@ class AuraTestHelper {
   // Flushes message loop.
   void RunAllPendingInMessageLoop();
 
-  Window* root_window() { return host_->window(); }
+  Window* root_window() { return host_ ? host_->window() : nullptr; }
   ui::EventSink* event_sink() { return host_->event_sink(); }
   WindowTreeHost* host() { return host_.get(); }
 
@@ -114,7 +115,10 @@ class AuraTestHelper {
   Mode mode_ = Mode::LOCAL;
   bool setup_called_;
   bool teardown_called_;
+  ui::ContextFactory* context_factory_to_restore_ = nullptr;
+  ui::ContextFactoryPrivate* context_factory_private_to_restore_ = nullptr;
   std::unique_ptr<TestWindowTreeClientSetup> window_tree_client_setup_;
+  Env::Mode env_mode_to_restore_ = Env::Mode::LOCAL;
   std::unique_ptr<aura::Env> env_;
   std::unique_ptr<wm::WMState> wm_state_;
   std::unique_ptr<WindowTreeHost> host_;

ui/aura/test/mus/window_tree_client_private.cc

@@ -38,7 +38,9 @@ WindowTreeHostMus* WindowTreeClientPrivate::CallWmNewDisplayAdded(
     const display::Display& display) {
   ui::mojom::WindowDataPtr root_data(ui::mojom::WindowData::New());
   root_data->parent_id = 0;
-  root_data->window_id = next_window_id_++;
+  // Windows representing displays are owned by mus, which is identified by
+  // non-zero high word.
+  root_data->window_id = next_window_id_++ | 0x00010000;
   root_data->visible = true;
   root_data->bounds = gfx::Rect(display.bounds().size());
   const bool parent_drawn = true;

ui/aura/test/test_screen.cc

@@ -55,7 +55,9 @@ WindowTreeHost* TestScreen::CreateHostForPrimaryDisplay() {
   // Makes sure InputMethod is default focused so that IME basics can work.
   host_->GetInputMethod()->OnFocus();
   host_->window()->AddObserver(this);
-  host_->InitHost();
+  // Other test code may have already initialized the compositor.
+  if (!host_->compositor()->root_layer())
+    host_->InitHost();
   host_->window()->Show();
   return host_;
 }

ui/views/mus/views_mus_test_suite.cc

@@ -24,8 +24,8 @@
 #include "services/service_manager/public/cpp/service_context.h"
 #include "services/ui/common/switches.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "ui/aura/env.h"
 #include "ui/aura/mus/window_tree_host_mus.h"
-#include "ui/aura/test/env_test_helper.h"
 #include "ui/aura/test/mus/input_method_mus_test_api.h"
 #include "ui/aura/window.h"
 #include "ui/compositor/test/fake_context_factory.h"
@@ -68,16 +68,13 @@ class PlatformTestHelperMus : public PlatformTestHelper {
  public:
   PlatformTestHelperMus(service_manager::Connector* connector,
                         const service_manager::Identity& identity) {
-    aura::test::EnvTestHelper().SetWindowTreeClient(nullptr);
     // It is necessary to recreate the MusClient for each test,
     // since a new MessageLoop is created for each test.
     mus_client_ = test::MusClientTestApi::Create(connector, identity);
     ViewsDelegate::GetInstance()->set_native_widget_factory(base::Bind(
         &PlatformTestHelperMus::CreateNativeWidget, base::Unretained(this)));
   }
-  ~PlatformTestHelperMus() override {
-    aura::test::EnvTestHelper().SetWindowTreeClient(nullptr);
-  }
+  ~PlatformTestHelperMus() override {}
 
   // PlatformTestHelper:
   void OnTestHelperCreated(ViewsTestHelper* helper) override {