Fix a bunch of mojo_public_*_unittests with the new EDK.

mojo/message_pump/message_pump_mojo_unittest.cc

Index: mojo/message_pump/message_pump_mojo_unittest.cc
diff --git a/mojo/message_pump/message_pump_mojo_unittest.cc b/mojo/message_pump/message_pump_mojo_unittest.cc
index 68e40df5eda5aeb6f25659772c426348a2cf2f10..0d370615c8dc48ccd860a612638756b09b9baac5 100644
--- a/mojo/message_pump/message_pump_mojo_unittest.cc
+++ b/mojo/message_pump/message_pump_mojo_unittest.cc
@@ -33,9 +33,29 @@ class CountingMojoHandler : public MessagePumpMojoHandler {
                    NULL,
                    MOJO_READ_MESSAGE_FLAG_NONE);
     ++success_count_;
+    if (success_count_ == success_callback_count_ &&
+        !success_callback_.is_null()) {
+      success_callback_.Run();
+      success_callback_.Reset();
+    }
   }
+
+  void set_success_callback(const base::Closure& callback,
+                            int success_count) {
+    success_callback_ = callback;
+    success_callback_count_ = success_count;
+  }
+
   void OnHandleError(const Handle& handle, MojoResult result) override {
     ++error_count_;
+    if (!error_callback_.is_null()) {
+      error_callback_.Run();
+      error_callback_.Reset();
+    }
+  }
+
+  void set_error_callback(const base::Closure& callback) {
+    error_callback_ = callback;
   }
 
   int success_count() { return success_count_; }
@@ -45,6 +65,11 @@ class CountingMojoHandler : public MessagePumpMojoHandler {
   int success_count_;
   int error_count_;
 
+  base::Closure error_callback_;
+  int success_callback_count_;
+
+  base::Closure success_callback_;
+
   DISALLOW_COPY_AND_ASSIGN(CountingMojoHandler);
 };
 
@@ -60,6 +85,8 @@ class CountingObserver : public MessagePumpMojo::Observer {
 TEST(MessagePumpMojo, RunUntilIdle) {
   base::MessageLoop message_loop(MessagePumpMojo::Create());
   CountingMojoHandler handler;
+  base::RunLoop run_loop;
+  handler.set_success_callback(run_loop.QuitClosure(), 2);
   MessagePipe handles;
   MessagePumpMojo::current()->AddHandler(&handler,
                                          handles.handle0.get(),
@@ -69,8 +96,11 @@ TEST(MessagePumpMojo, RunUntilIdle) {
       handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
   WriteMessageRaw(
       handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
-  base::RunLoop run_loop;
-  run_loop.RunUntilIdle();
+  MojoHandleSignalsState hss;
+  ASSERT_EQ(MOJO_RESULT_OK,
+            MojoWait(handles.handle0.get().value(), MOJO_HANDLE_SIGNAL_READABLE,
+                      MOJO_DEADLINE_INDEFINITE, &hss));
+  run_loop.Run();
   EXPECT_EQ(2, handler.success_count());
 }
 
@@ -81,6 +111,8 @@ TEST(MessagePumpMojo, Observer) {
   MessagePumpMojo::current()->AddObserver(&observer);
 
   CountingMojoHandler handler;
+  base::RunLoop run_loop;
+  handler.set_success_callback(run_loop.QuitClosure(), 1);
   MessagePipe handles;
   MessagePumpMojo::current()->AddHandler(&handler,
                                          handles.handle0.get(),
@@ -88,17 +120,25 @@ TEST(MessagePumpMojo, Observer) {
                                          base::TimeTicks());
   WriteMessageRaw(
       handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
-  base::RunLoop run_loop;
-  run_loop.RunUntilIdle();
+
+  MojoHandleSignalsState hss;
+  ASSERT_EQ(MOJO_RESULT_OK,
+            MojoWait(handles.handle0.get().value(), MOJO_HANDLE_SIGNAL_READABLE,
+                      MOJO_DEADLINE_INDEFINITE, &hss));
+  run_loop.Run();
   EXPECT_EQ(1, handler.success_count());
   EXPECT_EQ(1, observer.will_signal_handler_count);
   EXPECT_EQ(1, observer.did_signal_handler_count);
   MessagePumpMojo::current()->RemoveObserver(&observer);
 
+  base::RunLoop run_loop2;
+  handler.set_success_callback(run_loop2.QuitClosure(), 2);
   WriteMessageRaw(
       handles.handle1.get(), NULL, 0, NULL, 0, MOJO_WRITE_MESSAGE_FLAG_NONE);
-  base::RunLoop run_loop2;
-  run_loop2.RunUntilIdle();
+  ASSERT_EQ(MOJO_RESULT_OK,
+            MojoWait(handles.handle0.get().value(), MOJO_HANDLE_SIGNAL_READABLE,
+                      MOJO_DEADLINE_INDEFINITE, &hss));
+  run_loop2.Run();
   EXPECT_EQ(2, handler.success_count());
   EXPECT_EQ(1, observer.will_signal_handler_count);
   EXPECT_EQ(1, observer.did_signal_handler_count);
@@ -107,16 +147,15 @@ TEST(MessagePumpMojo, Observer) {
 TEST(MessagePumpMojo, UnregisterAfterDeadline) {
   base::MessageLoop message_loop(MessagePumpMojo::Create());
   CountingMojoHandler handler;
+  base::RunLoop run_loop;
+  handler.set_error_callback(run_loop.QuitClosure());
   MessagePipe handles;
   MessagePumpMojo::current()->AddHandler(
       &handler,
       handles.handle0.get(),
       MOJO_HANDLE_SIGNAL_READABLE,
       base::TimeTicks::Now() - base::TimeDelta::FromSeconds(1));
-  for (int i = 0; i < 2; ++i) {
-    base::RunLoop run_loop;
-    run_loop.RunUntilIdle();
-  }
+  run_loop.Run();
   EXPECT_EQ(1, handler.error_count());
 }