Connect to DesktopEnvironment's stubs only when audio/video schedulers are about to be created.

remoting/host/client_session_unittest.cc

Index: remoting/host/client_session_unittest.cc
diff --git a/remoting/host/client_session_unittest.cc b/remoting/host/client_session_unittest.cc
index 050304cff4ba6a2df6c9480dea7957fd2f221d1f..32cc8395217925bb81dfd774b543b148c7b51337 100644
--- a/remoting/host/client_session_unittest.cc
+++ b/remoting/host/client_session_unittest.cc
@@ -6,6 +6,7 @@
 #include "remoting/base/auto_thread_task_runner.h"
 #include "remoting/base/constants.h"
 #include "remoting/capturer/video_capturer_mock_objects.h"
+#include "remoting/capturer/video_frame_capturer_fake.h"
 #include "remoting/host/audio_capturer.h"
 #include "remoting/host/client_session.h"
 #include "remoting/host/desktop_environment.h"
@@ -26,107 +27,57 @@ using protocol::SessionConfig;
 using testing::_;
 using testing::AnyNumber;
 using testing::DeleteArg;
+using testing::DoAll;
 using testing::Expectation;
 using testing::InSequence;
 using testing::Return;
 using testing::ReturnRef;
 
+namespace {
+
+ACTION_P2(InjectClipboardEvent, connection, event) {
+  connection->clipboard_stub()->InjectClipboardEvent(event);
+}
+
+ACTION_P2(InjectKeyEvent, connection, event) {
+  connection->input_stub()->InjectKeyEvent(event);
+}
+
+ACTION_P2(InjectMouseEvent, connection, event) {
+  connection->input_stub()->InjectMouseEvent(event);
+}
+
+ACTION_P2(LocalMouseMoved, client_session, event) {
+  client_session->LocalMouseMoved(SkIPoint::Make(event.x(), event.y()));
+}
+
+}  // namespace
+
 class ClientSessionTest : public testing::Test {
  public:
   ClientSessionTest() : event_executor_(NULL) {}
 
-  virtual void SetUp() OVERRIDE {
-    ui_task_runner_ = new AutoThreadTaskRunner(
-        message_loop_.message_loop_proxy(),
-        base::Bind(&ClientSessionTest::QuitMainMessageLoop,
-                   base::Unretained(this)));
-
-    client_jid_ = "user@domain/rest-of-jid";
-
-    desktop_environment_factory_.reset(new MockDesktopEnvironmentFactory());
-    EXPECT_CALL(*desktop_environment_factory_, CreatePtr())
-        .Times(AnyNumber())
-        .WillRepeatedly(Invoke(this,
-                               &ClientSessionTest::CreateDesktopEnvironment));
-
-    // Set up a large default screen size that won't affect most tests.
-    screen_size_.set(1000, 1000);
-
-    session_config_ = SessionConfig::ForTest();
-
-    // Mock protocol::Session APIs called directly by ClientSession.
-    protocol::MockSession* session = new MockSession();
-    EXPECT_CALL(*session, config()).WillRepeatedly(ReturnRef(session_config_));
-    EXPECT_CALL(*session, jid()).WillRepeatedly(ReturnRef(client_jid_));
-    EXPECT_CALL(*session, SetEventHandler(_));
-
-    // Mock protocol::ConnectionToClient APIs called directly by ClientSession.
-    // HostStub is not touched by ClientSession, so we can safely pass NULL.
-    scoped_ptr<MockConnectionToClient> connection(
-        new MockConnectionToClient(session, NULL));
-    EXPECT_CALL(*connection, session()).WillRepeatedly(Return(session));
-    EXPECT_CALL(*connection, client_stub())
-        .WillRepeatedly(Return(&client_stub_));
-    EXPECT_CALL(*connection, video_stub()).WillRepeatedly(Return(&video_stub_));
-    EXPECT_CALL(*connection, Disconnect());
-    connection_ = connection.get();
-
-    client_session_ = new ClientSession(
-        &session_event_handler_,
-        ui_task_runner_, // Audio thread.
-        ui_task_runner_, // Capture thread.
-        ui_task_runner_, // Encode thread.
-        ui_task_runner_, // Network thread.
-        connection.PassAs<protocol::ConnectionToClient>(),
-        desktop_environment_factory_.get(),
-        base::TimeDelta());
-  }
+  virtual void SetUp() OVERRIDE;
+  virtual void TearDown() OVERRIDE;
 
-  virtual void TearDown() OVERRIDE {
-    // MockClientSessionEventHandler won't trigger Stop, so fake it.
-    client_session_->Stop(base::Bind(
-        &ClientSessionTest::OnClientStopped, base::Unretained(this)));
+  // Disconnects the client session.
+  void DisconnectClientSession();
 
-    // Run message loop before destroying because the session is destroyed
-    // asynchronously.
-    ui_task_runner_ = NULL;
-    message_loop_.Run();
-
-    // Verify that the client session has been stopped.
-    EXPECT_TRUE(client_session_.get() == NULL);
-  }
+  // Stops the client session.
+  void StopClientSession();

[Wez, 2013/01/08 18:34:00]

nit: Update this comment to explain what it means to stop the client session;
does this call cause ClientSession to tear-down, for example?

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

 
  protected:
-  DesktopEnvironment* CreateDesktopEnvironment() {
-    MockVideoFrameCapturer* capturer = new MockVideoFrameCapturer();
-    EXPECT_CALL(*capturer, Start(_));
-    EXPECT_CALL(*capturer, Stop());
-    EXPECT_CALL(*capturer, InvalidateRegion(_)).Times(AnyNumber());
-    EXPECT_CALL(*capturer, CaptureFrame()).Times(AnyNumber());
-    EXPECT_CALL(*capturer, size_most_recent())
-        .WillRepeatedly(ReturnRef(screen_size_));
-
-    EXPECT_TRUE(!event_executor_);
-    event_executor_ = new MockEventExecutor();
-    return new DesktopEnvironment(scoped_ptr<AudioCapturer>(NULL),
-                                  scoped_ptr<EventExecutor>(event_executor_),
-                                  scoped_ptr<VideoFrameCapturer>(capturer));
-  }
+  // Mocks DesktopEnvironmentFactory::Create();

[Wez, 2013/01/08 18:34:00]

nit: Creates a DesktopEnvironment with a fake VideoFrameCapturer, to mock
DesktopEnvironmentFactory::Create().

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+  DesktopEnvironment* CreateDesktopEnvironment();
 
-  void DisconnectClientSession() {
-    client_session_->Disconnect();
-    // MockSession won't trigger OnConnectionClosed, so fake it.
-    client_session_->OnConnectionClosed(client_session_->connection(),
-                                        protocol::OK);
-  }
+  // Connects the client session.

[Wez, 2013/01/08 18:34:00]

nit: To what? e.g. does it actually attach a fake ConnectionToClient, or
something else?

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+  void ConnectClientSession();
 
-  void QuitMainMessageLoop() {
-    message_loop_.PostTask(FROM_HERE, MessageLoop::QuitClosure());
-  }
+  // Quit the message loop.

[Wez, 2013/01/08 18:34:00]

typo: Quit -> Quits

nit: What affect does quitting the MessageLoop have? e.g. does it quit the test,
or just quit a mid-test run of the loop, or something else?

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+  void QuitMainMessageLoop();
 
-  void OnClientStopped() {
-    client_session_ = NULL;
-  }
+  // Invoked when the client session is fully stopped.

[Wez, 2013/01/08 18:34:00]

nit: What does it do? It seems to just release the ClientSession reference - is
it intended to indicate that the test can exit as well (via the
AutoThreadTaskRunner)?

e.g. Should it be:
// Releases the ClientSession when it has been fully stopped, allowing the
MessageLoop to quit.

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+  void OnClientStopped();
 
   // Message loop passed to |client_session_| to perform all functions on.
   MessageLoop message_loop_;
@@ -138,9 +89,6 @@ class ClientSessionTest : public testing::Test {
   // ClientSession::EventHandler mock for use in tests.
   MockClientSessionEventHandler session_event_handler_;
 
-  // Screen size that the fake VideoFrameCapturer should report.
-  SkISize screen_size_;
-
   // Storage for values to be returned by the protocol::Session mock.
   SessionConfig session_config_;
   std::string client_jid_;
@@ -159,6 +107,92 @@ class ClientSessionTest : public testing::Test {
   scoped_ptr<MockDesktopEnvironmentFactory> desktop_environment_factory_;
 };
 
+void ClientSessionTest::SetUp() {
+  ui_task_runner_ = new AutoThreadTaskRunner(

[Wez, 2013/01/08 18:34:00]

nit: Add some basic comments, e.g. "Arrange to run |message_loop_| until no
components depend on it" for this one.

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+      message_loop_.message_loop_proxy(),
+      base::Bind(&ClientSessionTest::QuitMainMessageLoop,
+                 base::Unretained(this)));
+
+  client_jid_ = "user@domain/rest-of-jid";

[Wez, 2013/01/08 18:34:00]

This doesn't ever seem to be set to any other value; can it be replaced with a
constant?

[alexeypa (please no reviews), 2013/01/08 20:00:14]

It has to be a member variable. It is not a POD type, so it cannot be global.
The mocked method returns a reference, so it has to be backed up by something.

[Wez, 2013/01/08 21:02:42]

Ah, OK.  Could it be a const member, initialized in the constructor?

+
+  desktop_environment_factory_.reset(new MockDesktopEnvironmentFactory());
+  EXPECT_CALL(*desktop_environment_factory_, CreatePtr())
+      .Times(AnyNumber())
+      .WillRepeatedly(Invoke(this,
+                             &ClientSessionTest::CreateDesktopEnvironment));
+
+  session_config_ = SessionConfig::ForTest();
+
+  // Mock protocol::Session APIs called directly by ClientSession.
+  protocol::MockSession* session = new MockSession();
+  EXPECT_CALL(*session, config()).WillRepeatedly(ReturnRef(session_config_));
+  EXPECT_CALL(*session, jid()).WillRepeatedly(ReturnRef(client_jid_));
+  EXPECT_CALL(*session, SetEventHandler(_));
+
+  // Mock protocol::ConnectionToClient APIs called directly by ClientSession.
+  // HostStub is not touched by ClientSession, so we can safely pass NULL.
+  scoped_ptr<MockConnectionToClient> connection(
+      new MockConnectionToClient(session, NULL));
+  EXPECT_CALL(*connection, session()).WillRepeatedly(Return(session));
+  EXPECT_CALL(*connection, client_stub())
+      .WillRepeatedly(Return(&client_stub_));
+  EXPECT_CALL(*connection, video_stub()).WillRepeatedly(Return(&video_stub_));
+  EXPECT_CALL(*connection, Disconnect());
+  connection_ = connection.get();
+
+  client_session_ = new ClientSession(
+      &session_event_handler_,
+      ui_task_runner_, // Audio thread.
+      ui_task_runner_, // Capture thread.
+      ui_task_runner_, // Encode thread.
+      ui_task_runner_, // Network thread.
+      connection.PassAs<protocol::ConnectionToClient>(),
+      desktop_environment_factory_.get(),
+      base::TimeDelta());
+  ui_task_runner_ = NULL;

[Wez, 2013/01/08 18:34:00]

Looks like ui_task_runner_ doesn't need to be a member of ClientSessionTest?

Or, you might move ui_task_runner_ = NULL into TearDown(), and put the
|message_loop_.Run()| call there.

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

[alexeypa (please no reviews), 2013/01/08 21:30:24]

Done.

+}
+
+void ClientSessionTest::TearDown() {
+  // Verify that the client session has been stopped.
+  EXPECT_TRUE(client_session_.get() == NULL);
+}
+
+void ClientSessionTest::DisconnectClientSession() {
+  client_session_->Disconnect();
+  // MockSession won't trigger OnConnectionClosed, so fake it.
+  client_session_->OnConnectionClosed(client_session_->connection(),
+                                      protocol::OK);
+}
+
+void ClientSessionTest::StopClientSession() {
+  // MockClientSessionEventHandler won't trigger Stop, so fake it.
+  client_session_->Stop(base::Bind(
+      &ClientSessionTest::OnClientStopped, base::Unretained(this)));
+}
+
+DesktopEnvironment* ClientSessionTest::CreateDesktopEnvironment() {
+  scoped_ptr<VideoFrameCapturer> video_capturer(new VideoFrameCapturerFake());
+
+  EXPECT_TRUE(!event_executor_);
+  event_executor_ = new MockEventExecutor();
+  return new DesktopEnvironment(scoped_ptr<AudioCapturer>(NULL),
+                                scoped_ptr<EventExecutor>(event_executor_),
+                                video_capturer.Pass());
+}
+
+void ClientSessionTest::ConnectClientSession() {
+  client_session_->OnConnectionAuthenticated(client_session_->connection());
+  client_session_->OnConnectionChannelsConnected(client_session_->connection());
+}
+
+void ClientSessionTest::QuitMainMessageLoop() {
+  message_loop_.PostTask(FROM_HERE, MessageLoop::QuitClosure());
+}
+
+void ClientSessionTest::OnClientStopped() {
+  client_session_ = NULL;
+}
+
 MATCHER_P2(EqualsClipboardEvent, m, d, "") {
   return (strcmp(arg.mime_type().c_str(), m) == 0 &&
       memcmp(arg.data().data(), d, arg.data().size()) == 0);
@@ -181,6 +215,15 @@ TEST_F(ClientSessionTest, ClipboardStubFilter) {
   EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(*event_executor_, StartPtr(_));
   EXPECT_CALL(session_event_handler_, OnSessionChannelsConnected(_));
+  EXPECT_CALL(video_stub_, ProcessVideoPacketPtr(_, _))

[Wez, 2013/01/08 18:34:00]

nit: Consider adding a comment to explain that this expectation uses the arrival
of the first video packet as a proxy for the transport streams having been
connected such that injecting input should work.

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+      .WillOnce(DoAll(
+          // This event should get through to the clipboard stub.
+          InjectClipboardEvent(connection_, clipboard_event2),
+          InvokeWithoutArgs(this, &ClientSessionTest::DisconnectClientSession),
+          // This event should not get through to the clipboard stub,
+          // because the client has disconnected.
+          InjectClipboardEvent(connection_, clipboard_event3),
+          InvokeWithoutArgs(this, &ClientSessionTest::StopClientSession)));
   EXPECT_CALL(*event_executor_, InjectClipboardEvent(EqualsClipboardEvent(
       kMimeTypeTextUtf8, "b")));
   EXPECT_CALL(session_event_handler_, OnSessionClosed(_));
@@ -188,14 +231,9 @@ TEST_F(ClientSessionTest, ClipboardStubFilter) {
   // This event should not get through to the clipboard stub,
   // because the client isn't authenticated yet.
   connection_->clipboard_stub()->InjectClipboardEvent(clipboard_event1);
-  client_session_->OnConnectionAuthenticated(client_session_->connection());
-  client_session_->OnConnectionChannelsConnected(client_session_->connection());
-  // This event should get through to the clipboard stub.
-  connection_->clipboard_stub()->InjectClipboardEvent(clipboard_event2);
-  DisconnectClientSession();
-  // This event should not get through to the clipboard stub,
-  // because the client has disconnected.
-  connection_->clipboard_stub()->InjectClipboardEvent(clipboard_event3);
+
+  ConnectClientSession();
+  message_loop_.Run();
 }
 
 MATCHER_P2(EqualsUsbEvent, usb_keycode, pressed, "") {
@@ -244,6 +282,18 @@ TEST_F(ClientSessionTest, InputStubFilter) {
   EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(*event_executor_, StartPtr(_));
   EXPECT_CALL(session_event_handler_, OnSessionChannelsConnected(_));
+  EXPECT_CALL(video_stub_, ProcessVideoPacketPtr(_, _))
+      .WillOnce(DoAll(
+          // These events should get through to the input stub.
+          InjectKeyEvent(connection_, key_event2_down),
+          InjectKeyEvent(connection_, key_event2_up),
+          InjectMouseEvent(connection_, mouse_event2),
+          InvokeWithoutArgs(this, &ClientSessionTest::DisconnectClientSession),
+          // These events should not get through to the input stub,
+          // because the client has disconnected.
+          InjectKeyEvent(connection_, key_event3),
+          InjectMouseEvent(connection_, mouse_event3),
+          InvokeWithoutArgs(this, &ClientSessionTest::StopClientSession)));
   EXPECT_CALL(*event_executor_, InjectKeyEvent(EqualsUsbEvent(2, true)));
   EXPECT_CALL(*event_executor_, InjectKeyEvent(EqualsUsbEvent(2, false)));
   EXPECT_CALL(*event_executor_, InjectMouseEvent(EqualsMouseEvent(200, 201)));
@@ -253,17 +303,9 @@ TEST_F(ClientSessionTest, InputStubFilter) {
   // because the client isn't authenticated yet.
   connection_->input_stub()->InjectKeyEvent(key_event1);
   connection_->input_stub()->InjectMouseEvent(mouse_event1);
-  client_session_->OnConnectionAuthenticated(client_session_->connection());
-  client_session_->OnConnectionChannelsConnected(client_session_->connection());
-  // These events should get through to the input stub.
-  connection_->input_stub()->InjectKeyEvent(key_event2_down);
-  connection_->input_stub()->InjectKeyEvent(key_event2_up);
-  connection_->input_stub()->InjectMouseEvent(mouse_event2);
-  DisconnectClientSession();
-  // These events should not get through to the input stub,
-  // because the client has disconnected.
-  connection_->input_stub()->InjectKeyEvent(key_event3);
-  connection_->input_stub()->InjectMouseEvent(mouse_event3);
+
+  ConnectClientSession();
+  message_loop_.Run();
 }
 
 TEST_F(ClientSessionTest, LocalInputTest) {
@@ -281,25 +323,26 @@ TEST_F(ClientSessionTest, LocalInputTest) {
   EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(*event_executor_, StartPtr(_));
   EXPECT_CALL(session_event_handler_, OnSessionChannelsConnected(_));
+  EXPECT_CALL(video_stub_, ProcessVideoPacketPtr(_, _))
+      .WillOnce(DoAll(
+          // This event should get through to the input stub.
+          InjectMouseEvent(connection_, mouse_event1),
+          // This one should too because the local event echoes the remote one.
+          LocalMouseMoved(client_session_.get(), mouse_event1),
+          InjectMouseEvent(connection_, mouse_event2),
+          // This one should not.
+          LocalMouseMoved(client_session_.get(), mouse_event1),
+          InjectMouseEvent(connection_, mouse_event3),
+          // TODO(jamiewalch): Verify that remote inputs are re-enabled
+          // eventually (via dependency injection, not sleep!)
+          InvokeWithoutArgs(this, &ClientSessionTest::DisconnectClientSession),
+          InvokeWithoutArgs(this, &ClientSessionTest::StopClientSession)));
   EXPECT_CALL(*event_executor_, InjectMouseEvent(EqualsMouseEvent(100, 101)));
   EXPECT_CALL(*event_executor_, InjectMouseEvent(EqualsMouseEvent(200, 201)));
   EXPECT_CALL(session_event_handler_, OnSessionClosed(_));
 
-  client_session_->OnConnectionAuthenticated(client_session_->connection());
-  client_session_->OnConnectionChannelsConnected(client_session_->connection());
-  // This event should get through to the input stub.
-  connection_->input_stub()->InjectMouseEvent(mouse_event1);
-  // This one should too because the local event echoes the remote one.
-  client_session_->LocalMouseMoved(SkIPoint::Make(mouse_event1.x(),
-                                                  mouse_event1.y()));
-  connection_->input_stub()->InjectMouseEvent(mouse_event2);
-  // This one should not.
-  client_session_->LocalMouseMoved(SkIPoint::Make(mouse_event1.x(),
-                                                  mouse_event1.y()));
-  connection_->input_stub()->InjectMouseEvent(mouse_event3);
-  // TODO(jamiewalch): Verify that remote inputs are re-enabled eventually
-  // (via dependency injection, not sleep!)
-  DisconnectClientSession();
+  ConnectClientSession();
+  message_loop_.Run();
 }
 
 TEST_F(ClientSessionTest, RestoreEventState) {
@@ -319,6 +362,13 @@ TEST_F(ClientSessionTest, RestoreEventState) {
   EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(*event_executor_, StartPtr(_));
   EXPECT_CALL(session_event_handler_, OnSessionChannelsConnected(_));
+  EXPECT_CALL(video_stub_, ProcessVideoPacketPtr(_, _))
+      .WillOnce(DoAll(
+          InjectKeyEvent(connection_, key1),
+          InjectKeyEvent(connection_, key2),
+          InjectMouseEvent(connection_, mousedown),
+          InvokeWithoutArgs(this, &ClientSessionTest::DisconnectClientSession),
+          InvokeWithoutArgs(this, &ClientSessionTest::StopClientSession)));
   EXPECT_CALL(*event_executor_, InjectKeyEvent(EqualsUsbEvent(1, true)));
   EXPECT_CALL(*event_executor_, InjectKeyEvent(EqualsUsbEvent(2, true)));
   EXPECT_CALL(*event_executor_, InjectMouseEvent(EqualsMouseButtonEvent(
@@ -329,48 +379,62 @@ TEST_F(ClientSessionTest, RestoreEventState) {
       protocol::MouseEvent::BUTTON_LEFT, false)));
   EXPECT_CALL(session_event_handler_, OnSessionClosed(_));
 
-  client_session_->OnConnectionAuthenticated(client_session_->connection());
-  client_session_->OnConnectionChannelsConnected(client_session_->connection());
-
-  connection_->input_stub()->InjectKeyEvent(key1);
-  connection_->input_stub()->InjectKeyEvent(key2);
-  connection_->input_stub()->InjectMouseEvent(mousedown);
-
-  DisconnectClientSession();
+  ConnectClientSession();
+  message_loop_.Run();
 }
 
 TEST_F(ClientSessionTest, ClampMouseEvents) {
-  screen_size_.set(200, 100);
-
   EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(*event_executor_, StartPtr(_));
   Expectation connected =
       EXPECT_CALL(session_event_handler_, OnSessionChannelsConnected(_));
   EXPECT_CALL(session_event_handler_, OnSessionClosed(_));
 
-  client_session_->OnConnectionAuthenticated(client_session_->connection());
-  client_session_->OnConnectionChannelsConnected(client_session_->connection());
-
   int input_x[3] = { -999, 100, 999 };
-  int expected_x[3] = { 0, 100, 199 };
+  int expected_x[3] = { 0, 100, VideoFrameCapturerFake::kWidth - 1 };
   int input_y[3] = { -999, 50, 999 };
-  int expected_y[3] = { 0, 50, 99 };
-
-  protocol::MouseEvent event;
+  int expected_y[3] = { 0, 50, VideoFrameCapturerFake::kHeight - 1 };
+
+  // Inject the 1st event after once a video packet has been received.

[Wez, 2013/01/08 18:34:00]

typo: "... after once ..."

[alexeypa (please no reviews), 2013/01/08 20:00:14]

Done.

+  protocol::MouseEvent injected_event;
+  injected_event.set_x(input_x[0]);
+  injected_event.set_y(input_y[0]);
+  connected =
+      EXPECT_CALL(video_stub_, ProcessVideoPacketPtr(_, _))
+          .After(connected)
+          .WillOnce(InjectMouseEvent(connection_, injected_event));
+
+  protocol::MouseEvent expected_event;
   for (int j = 0; j < 3; j++) {
     for (int i = 0; i < 3; i++) {
-      event.set_x(input_x[i]);
-      event.set_y(input_y[j]);
-      connected =
-          EXPECT_CALL(*event_executor_,
-                      InjectMouseEvent(EqualsMouseEvent(expected_x[i],
-                                                        expected_y[j])))
-              .After(connected);
-      connection_->input_stub()->InjectMouseEvent(event);
+      // Skip the first iteration since the 1st event has been injected already.
+      if (i > 0 || j > 0) {
+        injected_event.set_x(input_x[i]);
+        injected_event.set_y(input_y[j]);
+        connected =
+            EXPECT_CALL(*event_executor_,
+                        InjectMouseEvent(EqualsMouseEvent(expected_event.x(),
+                                                          expected_event.y())))
+                .After(connected)
+                .WillOnce(InjectMouseEvent(connection_, injected_event));
+      }
+
+      expected_event.set_x(expected_x[i]);
+      expected_event.set_y(expected_y[j]);
     }
   }
 
-  DisconnectClientSession();
+  // Shutdown the connection once the last event has been received.
+  EXPECT_CALL(*event_executor_,
+              InjectMouseEvent(EqualsMouseEvent(expected_event.x(),
+                                                expected_event.y())))
+      .After(connected)
+      .WillOnce(DoAll(
+          InvokeWithoutArgs(this, &ClientSessionTest::DisconnectClientSession),
+          InvokeWithoutArgs(this, &ClientSessionTest::StopClientSession)));
+
+  ConnectClientSession();
+  message_loop_.Run();
 }
 
 }  // namespace remoting