Support persistent session to make sure UIAutomation works with USB disconnection (#251)

* Support persistent session to make sure UIAutomation works with
USB disconnection
diff --git a/mobly/controllers/android_device_lib/snippet_client.py b/mobly/controllers/android_device_lib/snippet_client.py
index 3d85e40..f7f473b 100644
--- a/mobly/controllers/android_device_lib/snippet_client.py
+++ b/mobly/controllers/android_device_lib/snippet_client.py
@@ -24,15 +24,27 @@
     'com.google.android.mobly.snippet.SnippetRunner')
 
 # TODO(adorokhine): delete this in Mobly 1.6 when snippet v0 support is removed.
-_LAUNCH_CMD_V0 = ('am instrument -w -e action start -e port %s %s/' +
+_LAUNCH_CMD_V0 = ('%s am instrument -w -e action start -e port %s %s/' +
                   _INSTRUMENTATION_RUNNER_PACKAGE)
 
 _LAUNCH_CMD_V1 = (
-    'am instrument -w -e action start %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
+    '%s am instrument -w -e action start %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
 
 _STOP_CMD = (
     'am instrument -w -e action stop %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
 
+# Test that uses UiAutomation requires the shell session to be maintained while
+# test is in progress. However, this requirement does not hold for the test that
+# deals with device USB disconnection (Once device disconnects, the shell
+# session that started the instrument ends, and UiAutomation fails with error:
+# "UiAutomation not connected"). To keep the shell session and redirect
+# stdin/stdout/stderr, use "setsid" or "nohup" while launching the
+# instrumentation test. Because these commands may not be available in every
+# android system, try to use them only if exists.
+_SETSID_COMMAND = 'setsid'
+
+_NOHUP_COMMAND = 'nohup'
+
 # Maximum time to wait for a v0 snippet to start on the device (10 minutes).
 # TODO(adorokhine): delete this in Mobly 1.6 when snippet v0 support is removed.
 _APP_START_WAIT_TIME_V0 = 10 * 60
@@ -60,7 +72,7 @@
 
     def __init__(self, package, adb_proxy, log=logging.getLogger()):
         """Initializes a SnippetClient.
-  
+
         Args:
             package: (str) The package name of the apk where the snippets are
                      defined.
@@ -77,13 +89,14 @@
         """Overrides superclass. Launches a snippet app and connects to it."""
         self._check_app_installed()
 
+        persists_shell_cmd = self._get_persist_command()
         # Try launching the app with the v1 protocol. If that fails, fall back
         # to v0 for compatibility. Use info here so people know exactly what's
         # happening here, which is helpful since they need to create their own
         # instrumentations and manifest.
         self.log.info('Launching snippet apk %s with protocol v1',
                       self.package)
-        cmd = _LAUNCH_CMD_V1 % self.package
+        cmd = _LAUNCH_CMD_V1 % (persists_shell_cmd, self.package)
         start_time = time.time()
         self._proc = self._do_start_app(cmd)
 
@@ -106,7 +119,7 @@
             # Reuse the host port as the device port in v0 snippet. This isn't
             # safe in general, but the protocol is deprecated.
             self.device_port = self.host_port
-            cmd = _LAUNCH_CMD_V0 % (self.device_port, self.package)
+            cmd = _LAUNCH_CMD_V0 % (persists_shell_cmd, self.device_port, self.package)
             self._proc = self._do_start_app(cmd)
             self._connect_to_v0()
             self._launch_version = 'v0'
@@ -291,3 +304,17 @@
                 return line
             self.log.debug('Discarded line from instrumentation output: "%s"',
                            line)
+
+    def _get_persist_command(self):
+        """Check availability and return path of command if available."""
+        for command in [_SETSID_COMMAND, _NOHUP_COMMAND]:
+            try:
+                if command in self._adb.shell('which %s' % command):
+                    return command
+            except adb.AdbError:
+                continue
+        self.log.warning('No %s and %s commands available to launch instrument '
+                         'persistently, tests that depend on UiAutomator and '
+                         'at the same time performs USB disconnection may fail',
+                         _SETSID_COMMAND, _NOHUP_COMMAND)
+        return ''
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 010064c..beb9262 100755
--- a/tests/mobly/controllers/android_device_lib/snippet_client_test.py
+++ b/tests/mobly/controllers/android_device_lib/snippet_client_test.py
@@ -18,6 +18,7 @@
 import mock
 from future.tests.base import unittest
 
+from mobly.controllers.android_device_lib import adb
 from mobly.controllers.android_device_lib import jsonrpc_client_base
 from mobly.controllers.android_device_lib import snippet_client
 from tests.lib import jsonrpc_client_test_base
@@ -51,6 +52,8 @@
             return bytes('instrumentation:{p}/{r} (target={p})'.format(
                 p=MOCK_PACKAGE_NAME,
                 r=snippet_client._INSTRUMENTATION_RUNNER_PACKAGE), 'utf-8')
+        elif 'which' in params:
+            return ''
 
     def __getattr__(self, name):
         """All calls to the none-existent functions in adb proxy would
@@ -175,6 +178,73 @@
         client.start_app_and_connect()
         self.assertEqual(123, client.device_port)
 
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
+                'SnippetClient._do_start_app')
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
+                'SnippetClient._check_app_installed')
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
+                'SnippetClient._read_protocol_line')
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
+                'SnippetClient._connect_to_v1')
+    @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
+                'utils.get_available_host_port')
+    def test_snippet_start_app_and_connect_v1_persistent_session(
+            self, mock_get_port, mock_connect_to_v1, mock_read_protocol_line,
+            mock_check_app_installed, mock_do_start_app):
+
+        def _mocked_shell(arg):
+            if 'setsid' in arg:
+                raise adb.AdbError('cmd', 'stdout', 'stderr', 'ret_code')
+            else:
+                return 'nohup'
+
+        mock_get_port.return_value = 123
+        mock_read_protocol_line.side_effect = [
+            'SNIPPET START, PROTOCOL 1 234',
+            'SNIPPET SERVING, PORT 1234',
+            'SNIPPET START, PROTOCOL 1 234',
+            'SNIPPET SERVING, PORT 1234',
+            'SNIPPET START, PROTOCOL 1 234',
+            'SNIPPET SERVING, PORT 1234',
+        ]
+
+        # Test 'setsid' exists
+        client = self._make_client()
+        client._adb.shell = mock.Mock(return_value='setsid')
+        client.start_app_and_connect()
+        cmd_setsid = '%s am instrument -w -e action start %s/%s' % (
+            snippet_client._SETSID_COMMAND,
+            MOCK_PACKAGE_NAME,
+            snippet_client._INSTRUMENTATION_RUNNER_PACKAGE)
+        mock_do_start_app.assert_has_calls(mock.call(cmd_setsid))
+
+        # Test 'setsid' does not exist, but 'nohup' exsits
+        client = self._make_client()
+        client._adb.shell = _mocked_shell
+        client.start_app_and_connect()
+        cmd_nohup = '%s am instrument -w -e action start %s/%s' % (
+            snippet_client._NOHUP_COMMAND,
+            MOCK_PACKAGE_NAME,
+            snippet_client._INSTRUMENTATION_RUNNER_PACKAGE)
+        mock_do_start_app.assert_has_calls([
+            mock.call(cmd_setsid),
+            mock.call(cmd_nohup)
+        ])
+
+        # Test both 'setsid' and 'nohup' do not exist
+        client._adb.shell = mock.Mock(
+            side_effect=adb.AdbError('cmd', 'stdout', 'stderr', 'ret_code'))
+        client = self._make_client()
+        client.start_app_and_connect()
+        cmd_not_persist = ' am instrument -w -e action start %s/%s' % (
+            MOCK_PACKAGE_NAME,
+            snippet_client._INSTRUMENTATION_RUNNER_PACKAGE)
+        mock_do_start_app.assert_has_calls([
+            mock.call(cmd_setsid),
+            mock.call(cmd_nohup),
+            mock.call(cmd_not_persist)
+        ])
+
     @mock.patch('socket.create_connection')
     @mock.patch('mobly.controllers.android_device_lib.snippet_client.'
                 'utils.start_standing_subprocess')