diff --git a/ui/aura/env.cc b/ui/aura/env.cc
index c4dc11d1569b22..d4b5cdfe43d0d7 100644
--- a/ui/aura/env.cc
+++ b/ui/aura/env.cc
@@ -11,6 +11,8 @@
 #include "ui/aura/client/focus_client.h"
 #include "ui/aura/env_observer.h"
 #include "ui/aura/input_state_lookup.h"
+#include "ui/aura/mus/window_port_mus.h"
+#include "ui/aura/mus/window_tree_client.h"
 #include "ui/aura/window.h"
 #include "ui/aura/window_observer.h"
 #include "ui/aura/window_port_local.h"
@@ -78,10 +80,9 @@ Env::~Env() {
 }
 
 // static
-std::unique_ptr<Env> Env::CreateInstance(
-    const WindowPortFactory& window_port_factory) {
+std::unique_ptr<Env> Env::CreateInstance(Mode mode) {
   DCHECK(!lazy_tls_ptr.Pointer()->Get());
-  std::unique_ptr<Env> env(new Env(window_port_factory));
+  std::unique_ptr<Env> env(new Env(mode));
   env->Init();
   return env;
 }
@@ -100,9 +101,13 @@ Env* Env::GetInstanceDontCreate() {
 }
 
 std::unique_ptr<WindowPort> Env::CreateWindowPort(Window* window) {
-  if (window_port_factory_.is_null())
+  if (mode_ == Mode::LOCAL)
     return base::MakeUnique<WindowPortLocal>(window);
-  return window_port_factory_.Run(window);
+
+  DCHECK(window_tree_client_);
+  // Use LOCAL as all other cases are created by WindowTreeClient explicitly.
+  return base::MakeUnique<WindowPortMus>(window_tree_client_,
+                                         WindowMusType::LOCAL);
 }
 
 void Env::AddObserver(EnvObserver* observer) {
@@ -118,6 +123,13 @@ bool Env::IsMouseButtonDown() const {
       mouse_button_flags_ != 0;
 }
 
+void Env::SetWindowTreeClient(WindowTreeClient* window_tree_client) {
+  // The WindowTreeClient should only be set once. Test code may need to change
+  // the value after the fact, to do that use EnvTestHelper.
+  DCHECK(!window_tree_client_);
+  window_tree_client_ = window_tree_client;
+}
+
 void Env::SetActiveFocusClient(client::FocusClient* focus_client,
                                Window* focus_client_root) {
   if (focus_client == active_focus_client_ &&
@@ -139,8 +151,8 @@ void Env::SetActiveFocusClient(client::FocusClient* focus_client,
 ////////////////////////////////////////////////////////////////////////////////
 // Env, private:
 
-Env::Env(const WindowPortFactory& window_port_factory)
-    : window_port_factory_(window_port_factory),
+Env::Env(Mode mode)
+    : mode_(mode),
       mouse_button_flags_(0),
       is_touch_down_(false),
       input_state_lookup_(InputStateLookup::Create()),
diff --git a/ui/aura/env.h b/ui/aura/env.h
index acd289b4c580e8..274c9e74bc765e 100644
--- a/ui/aura/env.h
+++ b/ui/aura/env.h
@@ -7,7 +7,6 @@
 
 #include <memory>
 
-#include "base/callback.h"
 #include "base/macros.h"
 #include "base/observer_list.h"
 #include "base/supports_user_data.h"
@@ -34,17 +33,25 @@ class EnvObserver;
 class InputStateLookup;
 class Window;
 class WindowPort;
+class WindowTreeClient;
 class WindowTreeHost;
 
 // A singleton object that tracks general state within Aura.
 class AURA_EXPORT Env : public ui::EventTarget, public base::SupportsUserData {
  public:
+  enum class Mode {
+    // Classic aura.
+    LOCAL,
+
+    // Aura with a backend of mus.
+    MUS,
+  };
+
   ~Env() override;
 
-  using WindowPortFactory =
-      base::Callback<std::unique_ptr<WindowPort>(Window* window)>;
-  static std::unique_ptr<Env> CreateInstance(
-      const WindowPortFactory& window_port_factory = WindowPortFactory());
+  // NOTE: if you pass in Mode::MUS it is expected that you call
+  // SetWindowTreeClient() before any windows are created.
+  static std::unique_ptr<Env> CreateInstance(Mode mode = Mode::LOCAL);
   static Env* GetInstance();
   static Env* GetInstanceDontCreate();
 
@@ -78,6 +85,9 @@ class AURA_EXPORT Env : public ui::EventTarget, public base::SupportsUserData {
   }
   ui::ContextFactory* context_factory() { return context_factory_; }
 
+  // See CreateInstance() for description.
+  void SetWindowTreeClient(WindowTreeClient* window_tree_client);
+
   // Sets the active FocusClient and the window the FocusClient is associated
   // with. |window| is not necessarily the window that actually has focus.
   // |window| may be null, which indicates all windows share a FocusClient.
@@ -93,7 +103,7 @@ class AURA_EXPORT Env : public ui::EventTarget, public base::SupportsUserData {
   friend class Window;
   friend class WindowTreeHost;
 
-  explicit Env(const WindowPortFactory& window_port_factory);
+  explicit Env(Mode mode);
 
   void Init();
 
@@ -114,7 +124,10 @@ class AURA_EXPORT Env : public ui::EventTarget, public base::SupportsUserData {
   std::unique_ptr<ui::EventTargetIterator> GetChildIterator() const override;
   ui::EventTargeter* GetEventTargeter() override;
 
-  WindowPortFactory window_port_factory_;
+  // This is not const for tests, which may share Env across tests and so needs
+  // to reset the value.
+  Mode mode_;
+  WindowTreeClient* window_tree_client_ = nullptr;
 
   base::ObserverList<EnvObserver> observers_;
 
diff --git a/ui/aura/mus/focus_synchronizer.cc b/ui/aura/mus/focus_synchronizer.cc
index 8f69615ed114e6..b211315fc9ca0c 100644
--- a/ui/aura/mus/focus_synchronizer.cc
+++ b/ui/aura/mus/focus_synchronizer.cc
@@ -33,11 +33,13 @@ void FocusSynchronizer::SetFocusFromServer(WindowMus* window) {
   base::AutoReset<bool> focus_reset(&setting_focus_, true);
   base::AutoReset<WindowMus*> window_setting_focus_to_reset(
       &window_setting_focus_to_, window);
-  Env* env = aura::Env::GetInstance();
+  Env* env = Env::GetInstance();
   if (window) {
     Window* root = window->GetWindow()->GetRootWindow();
+    // The client should provide a focus client for all roots.
+    DCHECK(client::GetFocusClient(root));
     if (env->active_focus_client_root() != root)
-      env->SetActiveFocusClient(aura::client::GetFocusClient(root), root);
+      env->SetActiveFocusClient(client::GetFocusClient(root), root);
     window->GetWindow()->Focus();
   } else if (env->active_focus_client()) {
     env->active_focus_client()->FocusWindow(nullptr);
diff --git a/ui/aura/test/aura_test_helper.cc b/ui/aura/test/aura_test_helper.cc
index 501c3adf816238..3cd419e61aff7e 100644
--- a/ui/aura/test/aura_test_helper.cc
+++ b/ui/aura/test/aura_test_helper.cc
@@ -87,19 +87,21 @@ void AuraTestHelper::SetUp(ui::ContextFactory* context_factory) {
   // Needs to be before creating WindowTreeClient.
   focus_client_ = base::MakeUnique<TestFocusClient>();
   capture_client_ = base::MakeUnique<client::DefaultCaptureClient>();
-  Env::WindowPortFactory window_impl_factory =
-      base::Bind(&AuraTestHelper::CreateWindowPort, base::Unretained(this));
+  const Env::Mode env_mode =
+      (mode_ == Mode::LOCAL) ? Env::Mode::LOCAL : Env::Mode::MUS;
   if (mode_ == Mode::MUS_CREATE_WINDOW_TREE_CLIENT)
     InitWindowTreeClient();
   if (!Env::GetInstanceDontCreate())
-    env_ = Env::CreateInstance(window_impl_factory);
-  else
-    EnvTestHelper(Env::GetInstance()).SetWindowPortFactory(window_impl_factory);
+    env_ = Env::CreateInstance(env_mode);
+  EnvTestHelper env_helper;
+  // Always reset the mode. This really only matters for if Env was created
+  // above.
+  env_helper.SetMode(env_mode);
+  env_helper.SetWindowTreeClient(window_tree_client_);
   Env::GetInstance()->SetActiveFocusClient(focus_client_.get(), nullptr);
   Env::GetInstance()->set_context_factory(context_factory);
   // Unit tests generally don't want to query the system, rather use the state
   // from RootWindow.
-  EnvTestHelper env_helper(Env::GetInstance());
   env_helper.SetInputStateLookup(nullptr);
   env_helper.ResetEventState();
 
@@ -148,12 +150,10 @@ void AuraTestHelper::TearDown() {
 
   ui::ShutdownInputMethodForTesting();
 
-  if (env_) {
+  if (env_)
     env_.reset();
-  } else {
-    EnvTestHelper(Env::GetInstance())
-        .SetWindowPortFactory(Env::WindowPortFactory());
-  }
+  else
+    EnvTestHelper().SetWindowTreeClient(nullptr);
   wm_state_.reset();
 }
 
@@ -183,16 +183,5 @@ void AuraTestHelper::InitWindowTreeClient() {
   window_tree_client_ = window_tree_client_setup_->window_tree_client();
 }
 
-std::unique_ptr<WindowPort> AuraTestHelper::CreateWindowPort(Window* window) {
-  if (mode_ == Mode::LOCAL) {
-    std::unique_ptr<WindowPortLocal> window_port =
-        base::MakeUnique<WindowPortLocal>(window);
-    return std::move(window_port);
-  }
-  std::unique_ptr<WindowPortMus> window_port = base::MakeUnique<WindowPortMus>(
-      window_tree_client_, WindowMusType::LOCAL);
-  return std::move(window_port);
-}
-
 }  // namespace test
 }  // namespace aura
diff --git a/ui/aura/test/aura_test_helper.h b/ui/aura/test/aura_test_helper.h
index a1a5aa193b2c59..5390b879a21e29 100644
--- a/ui/aura/test/aura_test_helper.h
+++ b/ui/aura/test/aura_test_helper.h
@@ -31,7 +31,6 @@ class TestScreen;
 class TestWindowTree;
 class TestWindowTreeClientSetup;
 class Window;
-class WindowPort;
 class WindowManagerDelegate;
 class WindowTreeClient;
 class WindowTreeClientDelegate;
@@ -107,9 +106,6 @@ class AuraTestHelper {
   // Initializes a WindowTreeClient with a test WindowTree.
   void InitWindowTreeClient();
 
-  // Factory function for creating the appropriate WindowPort function.
-  std::unique_ptr<WindowPort> CreateWindowPort(Window* window);
-
   Mode mode_ = Mode::LOCAL;
   bool setup_called_;
   bool teardown_called_;
diff --git a/ui/aura/test/env_test_helper.h b/ui/aura/test/env_test_helper.h
index 2d7d7322ef3d7e..2402f3cff0f2ac 100644
--- a/ui/aura/test/env_test_helper.h
+++ b/ui/aura/test/env_test_helper.h
@@ -16,6 +16,7 @@ namespace test {
 
 class EnvTestHelper {
  public:
+  EnvTestHelper() : EnvTestHelper(Env::GetInstance()) {}
   explicit EnvTestHelper(Env* env) : env_(env) {}
   ~EnvTestHelper() {}
 
@@ -29,8 +30,12 @@ class EnvTestHelper {
     env_->is_touch_down_ = false;
   }
 
-  void SetWindowPortFactory(const Env::WindowPortFactory& factory) {
-    env_->window_port_factory_ = factory;
+  void SetMode(Env::Mode mode) { env_->mode_ = mode; }
+
+  // This circumvents the DCHECKs in Env::SetWindowTreeClient() and should
+  // only be used for tests where Env is long lived.
+  void SetWindowTreeClient(WindowTreeClient* window_tree_client) {
+    env_->window_tree_client_ = window_tree_client;
   }
 
  private:
diff --git a/ui/views/mus/BUILD.gn b/ui/views/mus/BUILD.gn
index f87f3fc50d1835..e06b14daf709dd 100644
--- a/ui/views/mus/BUILD.gn
+++ b/ui/views/mus/BUILD.gn
@@ -255,6 +255,7 @@ static_library("test_support_aura_mus") {
     "//services/ui/common:mus_common",
     "//testing/gtest",
     "//ui/aura",
+    "//ui/aura:test_support",
     "//ui/gl:test_support",
     "//ui/resources",
     "//ui/resources:ui_test_pak",
diff --git a/ui/views/mus/aura_init.cc b/ui/views/mus/aura_init.cc
index cc21e04d4ed227..38d7ae84a5fc3f 100644
--- a/ui/views/mus/aura_init.cc
+++ b/ui/views/mus/aura_init.cc
@@ -14,8 +14,6 @@
 #include "services/catalog/public/interfaces/constants.mojom.h"
 #include "services/service_manager/public/cpp/connector.h"
 #include "ui/aura/env.h"
-#include "ui/aura/mus/window_port_mus.h"
-#include "ui/aura/window_port_local.h"
 #include "ui/base/ime/input_method_initializer.h"
 #include "ui/base/material_design/material_design_controller.h"
 #include "ui/base/resource/resource_bundle.h"
@@ -57,8 +55,9 @@ AuraInit::AuraInit(service_manager::Connector* connector,
                    Mode mode)
     : resource_file_(resource_file),
       resource_file_200_(resource_file_200),
-      env_(aura::Env::CreateInstance(
-          base::Bind(&AuraInit::CreateWindowPort, base::Unretained(this)))),
+      env_(aura::Env::CreateInstance(mode == Mode::AURA_MUS
+                                         ? aura::Env::Mode::MUS
+                                         : aura::Env::Mode::LOCAL)),
       views_delegate_(new MusViewsDelegate) {
   if (mode == Mode::AURA_MUS) {
     mus_client_ =
@@ -118,17 +117,4 @@ void AuraInit::InitializeResources(service_manager::Connector* connector) {
         loader.TakeFile(resource_file_200_), ui::SCALE_FACTOR_200P);
 }
 
-std::unique_ptr<aura::WindowPort> AuraInit::CreateWindowPort(
-    aura::Window* window) {
-  if (mus_client_) {
-    std::unique_ptr<aura::WindowPortMus> window_port =
-        base::MakeUnique<aura::WindowPortMus>(mus_client_->window_tree_client(),
-                                              aura::WindowMusType::LOCAL);
-    return std::move(window_port);
-  }
-  std::unique_ptr<aura::WindowPortLocal> window_port =
-      base::MakeUnique<aura::WindowPortLocal>(window);
-  return std::move(window_port);
-}
-
 }  // namespace views
diff --git a/ui/views/mus/aura_init.h b/ui/views/mus/aura_init.h
index 9ef2f2d9d1c28e..c028a7e689cebc 100644
--- a/ui/views/mus/aura_init.h
+++ b/ui/views/mus/aura_init.h
@@ -16,7 +16,6 @@
 
 namespace aura {
 class Env;
-class WindowPort;
 }
 
 namespace base {
@@ -65,8 +64,6 @@ class VIEWS_MUS_EXPORT AuraInit {
  private:
   void InitializeResources(service_manager::Connector* connector);
 
-  std::unique_ptr<aura::WindowPort> CreateWindowPort(aura::Window* window);
-
 #if defined(OS_LINUX)
   sk_sp<font_service::FontLoader> font_loader_;
 #endif
diff --git a/ui/views/mus/mus_client.cc b/ui/views/mus/mus_client.cc
index 58b0b3eaaca536..53607b72383aba 100644
--- a/ui/views/mus/mus_client.cc
+++ b/ui/views/mus/mus_client.cc
@@ -96,6 +96,7 @@ MusClient::MusClient(service_manager::Connector* connector,
       compositor_context_factory_.get());
   window_tree_client_ =
       base::MakeUnique<aura::WindowTreeClient>(this, nullptr, nullptr);
+  aura::Env::GetInstance()->SetWindowTreeClient(window_tree_client_.get());
   window_tree_client_->ConnectViaWindowTreeFactory(connector_);
 
   // TODO: wire up PointerWatcherEventRouter. http://crbug.com/663526.
diff --git a/ui/views/mus/mus_client.h b/ui/views/mus/mus_client.h
index 879ccf73ae372e..f19685c90e775e 100644
--- a/ui/views/mus/mus_client.h
+++ b/ui/views/mus/mus_client.h
@@ -2,8 +2,8 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-#ifndef UI_VIEWS_MUS_MUS_CLIENT_INIT_H_
-#define UI_VIEWS_MUS_MUS_CLIENT_INIT_H_
+#ifndef UI_VIEWS_MUS_MUS_CLIENT_H_
+#define UI_VIEWS_MUS_MUS_CLIENT_H_
 
 #include <string>
 
@@ -128,4 +128,4 @@ class VIEWS_MUS_EXPORT MusClient
 
 }  // namespace views
 
-#endif  // UI_VIEWS_MUS_MUS_CLIENT_INIT_H_
+#endif  // UI_VIEWS_MUS_MUS_CLIENT_H_
diff --git a/ui/views/mus/views_aura_mus_test_suite.cc b/ui/views/mus/views_aura_mus_test_suite.cc
index 7badd395fe94a9..2278b949d03d5c 100644
--- a/ui/views/mus/views_aura_mus_test_suite.cc
+++ b/ui/views/mus/views_aura_mus_test_suite.cc
@@ -19,6 +19,7 @@
 #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/test/env_test_helper.h"
 #include "ui/views/mus/desktop_window_tree_host_mus.h"
 #include "ui/views/mus/mus_client.h"
 #include "ui/views/mus/test_utils.h"
@@ -55,11 +56,14 @@ 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);
   }
-  ~PlatformTestHelperMus() override {}
+  ~PlatformTestHelperMus() override {
+    aura::test::EnvTestHelper().SetWindowTreeClient(nullptr);
+  }
 
   // PlatformTestHelper:
   void OnTestHelperCreated(ViewsTestHelper* helper) override {