1 files changed, 40 insertions, 11 deletions

diff --git a/remoting/protocol/authenticator_test_base.cc b/remoting/protocol/authenticator_test_base.cc
index edff6d3..5ffe559 100644
--- a/remoting/protocol/authenticator_test_base.cc
+++ b/remoting/protocol/authenticator_test_base.cc
@@ -7,6 +7,8 @@
 #include "base/file_path.h"
 #include "base/file_util.h"
 #include "base/path_service.h"
+#include "base/test/test_timeouts.h"
+#include "base/timer.h"
 #include "crypto/rsa_private_key.h"
 #include "net/base/cert_test_util.h"
 #include "remoting/protocol/authenticator.h"
@@ -21,16 +23,25 @@ using testing::SaveArg;
 namespace remoting {
 namespace protocol {
 
-AuthenticatorTestBase::MockChannelDoneCallback::MockChannelDoneCallback() {
-}
-AuthenticatorTestBase::MockChannelDoneCallback::~MockChannelDoneCallback() {
-}
+namespace {
 
-AuthenticatorTestBase::AuthenticatorTestBase() {
-}
-AuthenticatorTestBase::~AuthenticatorTestBase() {
+ACTION_P(QuitThreadOnCounter, counter) {
+  --(*counter);
+  EXPECT_GE(*counter, 0);
+  if (*counter == 0)
+    MessageLoop::current()->Quit();
 }
 
+}  // namespace
+
+AuthenticatorTestBase::MockChannelDoneCallback::MockChannelDoneCallback() {}
+
+AuthenticatorTestBase::MockChannelDoneCallback::~MockChannelDoneCallback() {}
+
+AuthenticatorTestBase::AuthenticatorTestBase() {}
+
+AuthenticatorTestBase::~AuthenticatorTestBase() {}
+
 void AuthenticatorTestBase::SetUp() {
   FilePath certs_dir(net::GetTestCertsDirectory());
 
@@ -99,22 +110,40 @@ void AuthenticatorTestBase::RunChannelAuth(bool expected_fail) {
   net::StreamSocket* client_socket = NULL;
   net::StreamSocket* host_socket = NULL;
 
+  // Expect two callbacks to be called - the client callback and the host
+  // callback.
+  int callback_counter = 2;
+
   EXPECT_CALL(client_callback_, OnDone(net::OK, _))
-      .WillOnce(SaveArg<1>(&client_socket));
+      .WillOnce(DoAll(SaveArg<1>(&client_socket),
+                      QuitThreadOnCounter(&callback_counter)));
   if (expected_fail) {
-    EXPECT_CALL(host_callback_, OnDone(net::ERR_FAILED, NULL));
+    EXPECT_CALL(host_callback_, OnDone(net::ERR_FAILED, NULL))
+         .WillOnce(QuitThreadOnCounter(&callback_counter));
   } else {
     EXPECT_CALL(host_callback_, OnDone(net::OK, _))
-        .WillOnce(SaveArg<1>(&host_socket));
+        .WillOnce(DoAll(SaveArg<1>(&host_socket),
+                        QuitThreadOnCounter(&callback_counter)));
   }
 
-  message_loop_.RunAllPending();
+  // Ensure that .Run() does not run unbounded if the callbacks are never
+  // called.
+  base::Timer shutdown_timer(false, false);
+  shutdown_timer.Start(FROM_HERE, TestTimeouts::action_timeout(),
+                       MessageLoop::QuitClosure());
+  message_loop_.Run();
+  shutdown_timer.Stop();
 
   testing::Mock::VerifyAndClearExpectations(&client_callback_);
   testing::Mock::VerifyAndClearExpectations(&host_callback_);
 
   client_socket_.reset(client_socket);
   host_socket_.reset(host_socket);
+
+  if (!expected_fail) {
+    ASSERT_TRUE(client_socket_.get() != NULL);
+    ASSERT_TRUE(host_socket_.get() != NULL);
+  }
 }
 
 void AuthenticatorTestBase::OnHostConnected(