5 files changed, 19 insertions, 3 deletions
diff --git a/mojo/shell/tests/connect/connect_test_app.cc b/mojo/shell/tests/connect/connect_test_app.cc
index 4180763..342e14d 100644
--- a/mojo/shell/tests/connect/connect_test_app.cc
+++ b/mojo/shell/tests/connect/connect_test_app.cc
@@ -79,6 +79,12 @@ class ConnectTestApp : public ShellClient,
 
     return true;
   }
+  void ShellConnectionLost() override {
+    if (base::MessageLoop::current() &&
+        base::MessageLoop::current()->is_running()) {
+      base::MessageLoop::current()->QuitWhenIdle();
+    }
+  }
 
   // InterfaceFactory<test::mojom::ConnectTestService>:
   void Create(Connection* connection,
diff --git a/mojo/shell/tests/connect/connect_test_driver.cc b/mojo/shell/tests/connect/connect_test_driver.cc
index b1bdce4..ace64df 100644
--- a/mojo/shell/tests/connect/connect_test_driver.cc
+++ b/mojo/shell/tests/connect/connect_test_driver.cc
@@ -37,9 +37,8 @@ class Driver : public mojo::ShellClient,
     connection->AddInterface<ClientProcessTest>(this);
     return true;
   }
-  bool ShellConnectionLost() override {
-    // TODO(rockot): http://crbug.com/596621. Should be able to remove this
-    // override entirely.
+  void ShellConnectionLost() override {
+    // TODO: This should exit cleanly.
     _exit(1);
   }
 
diff --git a/mojo/shell/tests/connect/connect_test_package.cc b/mojo/shell/tests/connect/connect_test_package.cc
index 1458e65..c5a1a07 100644
--- a/mojo/shell/tests/connect/connect_test_package.cc
+++ b/mojo/shell/tests/connect/connect_test_package.cc
@@ -184,6 +184,12 @@ class ConnectTestShellClient
     connection->AddInterface<test::mojom::ConnectTestService>(this);
     return true;
   }
+  void ShellConnectionLost() override {
+    if (base::MessageLoop::current() &&
+        base::MessageLoop::current()->is_running()) {
+      base::MessageLoop::current()->QuitWhenIdle();
+    }
+  }
 
   // InterfaceFactory<mojom::ShellClientFactory>:
   void Create(Connection* connection,
diff --git a/mojo/shell/tests/lifecycle/app_client.cc b/mojo/shell/tests/lifecycle/app_client.cc
index 657d967..a2400ec 100644
--- a/mojo/shell/tests/lifecycle/app_client.cc
+++ b/mojo/shell/tests/lifecycle/app_client.cc
@@ -20,6 +20,10 @@ bool AppClient::AcceptConnection(mojo::Connection* connection) {
   return true;
 }
 
+void AppClient::ShellConnectionLost() {
+  GracefulQuit();
+}
+
 void AppClient::Create(mojo::Connection* connection,
                        LifecycleControlRequest request) {
   bindings_.AddBinding(this, std::move(request));
diff --git a/mojo/shell/tests/lifecycle/app_client.h b/mojo/shell/tests/lifecycle/app_client.h
index 5a7d5e6..d389f20 100644
--- a/mojo/shell/tests/lifecycle/app_client.h
+++ b/mojo/shell/tests/lifecycle/app_client.h
@@ -37,6 +37,7 @@ class AppClient : public ShellClient,
 
   // ShellClient:
   bool AcceptConnection(Connection* connection) override;
+  void ShellConnectionLost() override;
 
   // InterfaceFactory<LifecycleControl>:
   void Create(Connection* connection, LifecycleControlRequest request) override;