Fix the bug where event client is broken after reboot (#821)

mobly/controllers/android_device_lib/jsonrpc_client_base.py

@@ -253,13 +253,17 @@ def disconnect(self):
     `SnippetClient.restore_app_connection`.
     """
     try:
-      if self._conn:
-        self._conn.close()
-        self._conn = None
+      self.close_socket_connection()
     finally:
       # Always clear the host port as part of the disconnect step.
       self.clear_host_port()
 
+  def close_socket_connection(self):
+    """Closes the socket connection to the server."""
+    if self._conn:
+      self._conn.close()
+      self._conn = None
+
   def clear_host_port(self):
     """Stops the adb port forwarding of the host port used by this client.
     """

mobly/controllers/android_device_lib/snippet_client.py

@@ -251,6 +251,8 @@ def stop_app(self):
       raise errors.DeviceError(
           self._ad, 'Failed to stop existing apk. Unexpected output: %s' % out)
 
+    self._stop_event_client()
+
   def _start_event_client(self):
     """Overrides superclass."""
     event_client = SnippetClient(package=self.package, ad=self._ad)
@@ -259,6 +261,12 @@ def _start_event_client(self):
     event_client.connect(self.uid, jsonrpc_client_base.JsonRpcCommand.CONTINUE)
     return event_client
 
+  def _stop_event_client(self):
+    """Releases all the resources acquired in `_start_event_client`."""
+    if self._event_client:
+      self._event_client.close_socket_connection()
+      self._event_client = None
+
   def _restore_event_client(self):
     """Restores previously created event client."""
     if not self._event_client:

tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py

@@ -337,6 +337,22 @@ def test_rpc_truncated_logging_long_response(self, mock_create_connection):
         testing_rpc_response[:jsonrpc_client_base._MAX_RPC_RESP_LOGGING_LENGTH],
         resp_len - jsonrpc_client_base._MAX_RPC_RESP_LOGGING_LENGTH)
 
+  def test_close_scoket_connection(self):
+    client = FakeRpcClient()
+    mock_conn = mock.Mock()
+    client._conn = mock_conn
+
+    client.close_socket_connection()
+    mock_conn.close.assert_called_once()
+    self.assertIsNone(client._conn)
+
+  def test_close_scoket_connection_without_connection(self):
+    client = FakeRpcClient()
+    client._conn = None
+
+    client.close_socket_connection()
+    self.assertIsNone(client._conn)
+
 
 if __name__ == '__main__':
   unittest.main()

tests/mobly/controllers/android_device_lib/snippet_client_test.py

@@ -177,6 +177,56 @@ def test_snippet_stop_app_raises(self):
       client.stop_app()
     adb_proxy.forward.assert_called_once_with(['--remove', 'tcp:1'])
 
+  @mock.patch('socket.create_connection')
+  @mock.patch('mobly.utils.stop_standing_subprocess')
+  def test_snippet_stop_app_stops_event_client(self,
+                                               mock_stop_standing_subprocess,
+                                               mock_create_connection):
+    adb_proxy = mock.MagicMock()
+    adb_proxy.shell.return_value = b'OK (0 tests)'
+    client = self._make_client(adb_proxy)
+    event_client = snippet_client.SnippetClient(
+        package=MOCK_PACKAGE_NAME, ad=client._ad)
+    client._event_client = event_client
+    event_client_conn = mock.Mock()
+    event_client._conn = event_client_conn
+
+    client.stop_app()
+    self.assertFalse(client.is_alive)
+    event_client_conn.close.assert_called_once()
+    self.assertIsNone(client._event_client)
+    self.assertIsNone(event_client._conn)
+
+  @mock.patch('socket.create_connection')
+  @mock.patch('mobly.utils.stop_standing_subprocess')
+  def test_snippet_stop_app_stops_event_client_without_connection(
+      self, mock_stop_standing_subprocess, mock_create_connection):
+    adb_proxy = mock.MagicMock()
+    adb_proxy.shell.return_value = b'OK (0 tests)'
+    client = self._make_client(adb_proxy)
+    event_client = snippet_client.SnippetClient(
+        package=MOCK_PACKAGE_NAME, ad=client._ad)
+    client._event_client = event_client
+    event_client._conn = None
+
+    client.stop_app()
+    self.assertFalse(client.is_alive)
+    self.assertIsNone(client._event_client)
+    self.assertIsNone(event_client._conn)
+
+  @mock.patch('socket.create_connection')
+  @mock.patch('mobly.utils.stop_standing_subprocess')
+  def test_snippet_stop_app_without_event_client(
+      self, mock_stop_standing_subprocess, mock_create_connection):
+    adb_proxy = mock.MagicMock()
+    adb_proxy.shell.return_value = b'OK (0 tests)'
+    client = self._make_client(adb_proxy)
+    client._event_client = None
+
+    client.stop_app()
+    self.assertFalse(client.is_alive)
+    self.assertIsNone(client._event_client)
+
   @mock.patch('socket.create_connection')
   @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
               'utils.start_standing_subprocess')