From fb0a3969e87763e089b28261d3cf4f4c5d64e82f Mon Sep 17 00:00:00 2001 From: Minghao Li Date: Wed, 18 May 2022 16:07:32 +0800 Subject: [PATCH] Fix the bug where event client is broken after reboot (#821) --- .../android_device_lib/jsonrpc_client_base.py | 10 ++-- .../android_device_lib/snippet_client.py | 8 +++ .../jsonrpc_client_base_test.py | 16 ++++++ .../android_device_lib/snippet_client_test.py | 50 +++++++++++++++++++ 4 files changed, 81 insertions(+), 3 deletions(-) diff --git a/mobly/controllers/android_device_lib/jsonrpc_client_base.py b/mobly/controllers/android_device_lib/jsonrpc_client_base.py index 39266dca..7d7e3a12 100644 --- a/mobly/controllers/android_device_lib/jsonrpc_client_base.py +++ b/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. """ diff --git a/mobly/controllers/android_device_lib/snippet_client.py b/mobly/controllers/android_device_lib/snippet_client.py index a15b52c3..c0dbadc0 100644 --- a/mobly/controllers/android_device_lib/snippet_client.py +++ b/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: diff --git a/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py b/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py index 96b52cbe..4cbeb35e 100755 --- a/tests/mobly/controllers/android_device_lib/jsonrpc_client_base_test.py +++ b/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() diff --git a/tests/mobly/controllers/android_device_lib/snippet_client_test.py b/tests/mobly/controllers/android_device_lib/snippet_client_test.py index 6bb2f87c..bfa41d6f 100755 --- a/tests/mobly/controllers/android_device_lib/snippet_client_test.py +++ b/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')