Repository: google/mobly

Issue:
[Snippet Client] Backward compatibility mechanism with older versions of snippet lib
If we make updates to snippet lib and introduce new internal Rpcs or other protocol changes, the client needs to be able to figure out which version of the lib is used and behave accordingly.

Discussion:
xpconanfan: 



Review status: 0 of 10 files reviewed at latest revision, 6 unresolved discussions.

---

*[mobly/utils.py, line 191 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksmGNXOJuTxDvAq8CM:-KksmGNXOJuTxDvAq8CN:b-oqfifj) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/utils.py#L191)):*
> ```Python
>     for path in paths:
>         p = abs_path(path)
>         for dirPath, unused_subdirList, fileList in os.walk(p):
> ```

Simply use `_` instead of `unused_subdirList`?

---

*[mobly/utils.py, line 429 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksmQfoOeHMN6bqR-az:-KksmQfoOeHMN6bqR-b-:b8h2k1j) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/utils.py#L429)):*
> ```Python
>     that matches a given regex pattern.
> 
>     This function is specifically used to grep strings from AdbProxy's
> ```

If this is specific to `AdbProxy`, should it live there instead?

---

*[mobly/controllers/android_device_lib/jsonrpc_client_base.py, line 124 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksppqrXEQ8fimOOv0j:-KksppqsL48x2i169ytC:b-5fd9cq) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/jsonrpc_client_base.py#L124)):*
> ```Python
>     # Methods to be implemented by subclasses.
> 
>     def start_app(self):
> ```

Should this method's name be changed to something like `start_app_and_connect`?
(as long as the new name indicates that it also connects)

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 35 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksq6k8LPXjF8h3xhjB:-Kksq6k8LPXjF8h3xhjC:b-flbu2r) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L35)):*
> ```Python
>     'am instrument -w -e action stop %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
> 
> # Maximum time to wait for the app to start on the device (10 minutes).
> ```

Could we keep this in base class and assign `_APP_START_WAIT_TIME_V0` as an alias, so we don't have to dup this in two places.

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 50 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksnWRcgnlQ1Gwp-jor:-KksnWRcgnlQ1Gwp-jos:bgyt3m8) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L50)):*
> ```Python
> 
> 
> class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
> ```

Can we add some context on the v0 and v1 changes in the docstring here?

---

*[tools/sl4a_shell.py, line 67 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksp8Cgj7VeRpqF61_B:-Kksp8Cgj7VeRpqF61_C:bqebhii) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/tools/sl4a_shell.py#L67)):*
> ```Python
>         'Device serial to connect to (if more than one device is connected)')
>     args = parser.parse_args()
>     logging.basicConfig(level=logging.INFO)
> ```

why this change?

---


*Comments from [Reviewable](https://reviewable.io:443/reviews/google/mobly/216)*
<!-- Sent from Reviewable.io -->

adorokhine: 



Review status: 0 of 10 files reviewed at latest revision, 6 unresolved discussions.

---

*[mobly/utils.py, line 191 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksmGNXOJuTxDvAq8CM:-Kkt8twZgicerHD2OJMW:b-896fix) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/utils.py#L191)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

Simply use `_` instead of `unused_subdirList`?
</blockquote></details>

Done.

---

*[mobly/utils.py, line 429 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksmQfoOeHMN6bqR-az:-Kkt9M0sFzwLJ9oFQD7k:bpqwp91) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/utils.py#L429)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

If this is specific to `AdbProxy`, should it live there instead?
</blockquote></details>

Changed this text; it's not specific to AdbProxy.

---

*[mobly/controllers/android_device_lib/jsonrpc_client_base.py, line 124 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksppqrXEQ8fimOOv0j:-Kkt9t1A400G0mYJUpPk:b-896fix) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/jsonrpc_client_base.py#L124)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

Should this method's name be changed to something like `start_app_and_connect`?
(as long as the new name indicates that it also connects)
</blockquote></details>

Done.

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 35 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksq6k8LPXjF8h3xhjB:-Kkt9y5p6eip75Lgct2X:b-umu14f) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L35)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

Could we keep this in base class and assign `_APP_START_WAIT_TIME_V0` as an alias, so we don't have to dup this in two places.
</blockquote></details>

Actually it shouldn't be duped, because the start time of sl4a has nothing to do with the start time of snippet. I want to reduce the sl4a timeout again but I don't want to do that in this PR.

The one in snippet will soon be deleted because v1 snippets don't have a timeout value like this.

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 50 at r1](https://reviewable.io:443/reviews/google/mobly/216#-KksnWRcgnlQ1Gwp-jor:-KktAVbpYGDTTWN_-E_I:b-896fix) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L50)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

Can we add some context on the v0 and v1 changes in the docstring here?
</blockquote></details>

Done.

---

*[tools/sl4a_shell.py, line 67 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksp8Cgj7VeRpqF61_B:-KktAWQ1epwCDO8KN5OZ:bwqbmaq) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/tools/sl4a_shell.py#L67)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

why this change?
</blockquote></details>

The same line already exists in snippet_shell so this makes it consistent. It makes it easier to debug because you can just flip this value and rerun to get debug logs. (there's no file logging happening in the shells so there's no other way to view debug logs at the moment.)

---


*Comments from [Reviewable](https://reviewable.io:443/reviews/google/mobly/216)*
<!-- Sent from Reviewable.io -->

xpconanfan: 



Review status: 0 of 10 files reviewed at latest revision, 3 unresolved discussions.

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 35 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksq6k8LPXjF8h3xhjB:-KktGTB7IK6Eci8jgw2z:b-mhl3j7) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L35)):*
<details><summary><i>Previously, adorokhine (Alexander Dorokhine) wrote…</i></summary><blockquote>

Actually it shouldn't be duped, because the start time of sl4a has nothing to do with the start time of snippet. I want to reduce the sl4a timeout again but I don't want to do that in this PR.

The one in snippet will soon be deleted because v1 snippets don't have a timeout value like this.
</blockquote></details>

Should this comment mention the pending removal of this value?
Should the comment for the one in sl4a mention the coming reduction in that value?

---


*Comments from [Reviewable](https://reviewable.io:443/reviews/google/mobly/216)*
<!-- Sent from Reviewable.io -->

adorokhine: 



Review status: 0 of 10 files reviewed at latest revision, 3 unresolved discussions.

---

*[mobly/controllers/android_device_lib/snippet_client.py, line 35 at r1](https://reviewable.io:443/reviews/google/mobly/216#-Kksq6k8LPXjF8h3xhjB:-KktIKdGuZz_M4fjY_ic:b-896fix) ([raw file](https://github.com/google/mobly/blob/b3fac107ec424a0b8bce0dc1492db2921ef986da/mobly/controllers/android_device_lib/snippet_client.py#L35)):*
<details><summary><i>Previously, xpconanfan (Ang Li) wrote…</i></summary><blockquote>

Should this comment mention the pending removal of this value?
Should the comment for the one in sl4a mention the coming reduction in that value?
</blockquote></details>

Done.

---


*Comments from [Reviewable](https://reviewable.io:443/reviews/google/mobly/216)*
<!-- Sent from Reviewable.io -->

xpconanfan: <img class="emoji" title=":lgtm:" alt=":lgtm:" align="absmiddle" src="https://reviewable.io/lgtm.png" height="20" width="61"/>

---

Review status: 0 of 10 files reviewed at latest revision, all discussions resolved.

---



*Comments from [Reviewable](https://reviewable.io:443/reviews/google/mobly/216#-:-Kl0YT3OyOKHOb-QvWIa:bnfp4nl)*
<!-- Sent from Reviewable.io -->


Patch:
diff --git a/mobly/controllers/android_device.py b/mobly/controllers/android_device.py
index 3f8e209..33f7309 100644
--- a/mobly/controllers/android_device.py
+++ b/mobly/controllers/android_device.py
@@ -25,9 +25,7 @@ from mobly import logger as mobly_logger
 from mobly import signals
 from mobly import utils
 from mobly.controllers.android_device_lib import adb
-from mobly.controllers.android_device_lib import event_dispatcher
 from mobly.controllers.android_device_lib import fastboot
-from mobly.controllers.android_device_lib import jsonrpc_client_base
 from mobly.controllers.android_device_lib import sl4a_client
 from mobly.controllers.android_device_lib import snippet_client
 
@@ -462,7 +460,7 @@ class AndroidDevice(object):
         service_info['use_sl4a'] = self.sl4a is not None
         self._terminate_sl4a()
         for name, client in self._snippet_clients.items():
-            self._terminate_jsonrpc_client(client)
+            client.stop_app()
             delattr(self, name)
         self._snippet_clients = {}
         if self._adb_logcat_process:
@@ -643,13 +641,9 @@ class AndroidDevice(object):
                     self,
                     'Snippet package "%s" has already been loaded under name'
                     ' "%s".' % (package, client_name))
-        host_port = utils.get_available_host_port()
         client = snippet_client.SnippetClient(
-            package=package,
-            host_port=host_port,
-            adb_proxy=self.adb,
-            log=self.log)
-        self._start_jsonrpc_client(client)
+            package=package, adb_proxy=self.adb, log=self.log)
+        client.start_app_and_connect()
         self._snippet_clients[name] = client
         setattr(self, name, client)
 
@@ -662,48 +656,10 @@ class AndroidDevice(object):
         Creates an sl4a client (self.sl4a) with one connection, and one
         EventDispatcher obj (self.ed) with the other connection.
         """
-        host_port = utils.get_available_host_port()
-        self.sl4a = sl4a_client.Sl4aClient(
-            host_port=host_port, adb_proxy=self.adb)
-        self._start_jsonrpc_client(self.sl4a)
-
-        # Start an EventDispatcher for the current sl4a session
-        event_client = sl4a_client.Sl4aClient(
-            host_port=host_port, adb_proxy=self.adb)
-        event_client.connect(
-            uid=self.sl4a.uid, cmd=jsonrpc_client_base.JsonRpcCommand.CONTINUE)
-        self.ed = event_dispatcher.EventDispatcher(event_client)
-        self.ed.start()
-
-    def _start_jsonrpc_client(self, client):
-        """Create a connection to a jsonrpc server running on the device.
-
-        If the connection cannot be made, tries to restart it.
-        """
-        client.check_app_installed()
-        self.adb.forward(
-            ['tcp:%d' % client.host_port, 'tcp:%d' % client.device_port])
-        try:
-            client.connect()
-        except:
-            try:
-                client.stop_app()
-            except Exception as e:
-                self.log.warning(e)
-            client.start_app()
-            client.connect()
-
-    def _terminate_jsonrpc_client(self, client):
-        try:
-            client.closeSl4aSession()
-            client.close()
-            client.stop_app()
-        except:
-            self.log.exception('Failed to stop Rpc client for %s.',
-                               client.app_name)
-        finally:
-            # Always clean up the adb port
-            self.adb.forward(['--remove', 'tcp:%d' % client.host_port])
+        self.sl4a = sl4a_client.Sl4aClient(adb_proxy=self.adb, log=self.log)
+        self.sl4a.start_app_and_connect()
+        # Unpack the 'ed' attribute for compatibility.
+        self.ed = self.sl4a.ed
 
     def _is_timestamp_in_range(self, target, begin_time, end_time):
         low = mobly_logger.logline_timestamp_comparator(begin_time,
@@ -849,13 +805,8 @@ class AndroidDevice(object):
         the session. Clear corresponding droids and dispatchers from cache.
         """
         if self.sl4a:
-            self._terminate_jsonrpc_client(self.sl4a)
+            self.sl4a.stop_app()
             self.sl4a = None
-        if self.ed:
-            try:
-                self.ed.clean_up()
-            except:
-                self.log.exception('Failed to shutdown sl4a event dispatcher.')
             self.ed = None
 
     def run_iperf_client(self, server_host, extra_args=''):
diff --git a/mobly/controllers/android_device_lib/event_dispatcher.py b/mobly/controllers/android_device_lib/event_dispatcher.py
index a718ca4..5ebf4df 100644
--- a/mobly/controllers/android_device_lib/event_dispatcher.py
+++ b/mobly/controllers/android_device_lib/event_dispatcher.py
@@ -147,7 +147,10 @@ class EventDispatcher:
             return
         self.started = False
         self.clear_all_events()
-        self._sl4a.close()
+        # At this point, the sl4a apk is destroyed and nothing is listening on
+        # the socket. Avoid sending any sl4a commands; just clean up the socket
+        # and return.
+        self._sl4a.disconnect()
         self.poller.set_result("Done")
         # The polling thread is guaranteed to finish after a max of 60 seconds,
         # so we don't wait here.
diff --git a/mobly/controllers/android_device_lib/jsonrpc_client_base.py b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
index 9f96aaa..5468271 100644
--- a/mobly/controllers/android_device_lib/jsonrpc_client_base.py
+++ b/mobly/controllers/android_device_lib/jsonrpc_client_base.py
@@ -38,20 +38,11 @@ from builtins import str
 
 import json
 import logging
-import re
 import socket
 import threading
-import time
 
-from mobly.controllers.android_device_lib import adb
 from mobly.controllers.android_device_lib import callback_handler
 
-# Maximum time to wait for the app to start on the device (10 minutes).
-# TODO: This timeout is set high in order to allow for retries in start_app.
-# Decrease it when the call to connect() has the option for a quicker timeout
-# than the default _cmd() timeout.
-APP_START_WAIT_TIME = 10 * 60
-
 # UID of the 'unknown' jsonrpc session. Will cause creation of a new session.
 UNKNOWN_UID = -1
 
@@ -107,39 +98,44 @@ class JsonRpcClientBase(object):
         uid: (int) The uid of this session.
     """
 
-    def __init__(self,
-                 host_port,
-                 device_port,
-                 app_name,
-                 adb_proxy,
-                 log=logging.getLogger()):
+    def __init__(self, app_name, log=logging.getLogger()):
         """
         Args:
-            host_port: (int) The host port of this RPC client.
-            device_port: (int) The device port of this RPC client.
             app_name: (str) The user-visible name of the app being communicated
                       with.
-            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
+            log: (logging.Logger) logger to which to send log messages.
         """
-        self.host_port = host_port
-        self.device_port = device_port
+        self.host_port = None
+        self.device_port = None
         self.app_name = app_name
+        self.log = log
         self.uid = None
-        self._adb = adb_proxy
         self._client = None  # prevent close errors on connect failure
         self._conn = None
         self._counter = None
         self._lock = threading.Lock()
         self._event_client = None
-        self._log = log
 
     def __del__(self):
-        self.close()
+        self.disconnect()
 
     # Methods to be implemented by subclasses.
 
-    def _do_start_app(self):
-        """Starts the server app on the android device.
+    def start_app_and_connect(self):
+        """Starts the server app on the android device and connects to it.
+
+        After this, the self.host_port and self.device_port attributes must be
+        set.
+
+        Must be implemented by subclasses.
+
+        Raises:
+            AppStartError: When the app was not able to be started.
+        """
+        raise NotImplementedError()
+
+    def stop_app(self):
+        """Kills any running instance of the app.
 
         Must be implemented by subclasses.
         """
@@ -158,47 +154,8 @@ class JsonRpcClientBase(object):
         """
         raise NotImplementedError()
 
-    def stop_app(self):
-        """Kills any running instance of the app.
-
-        Must be implemented by subclasses.
-        """
-        raise NotImplementedError()
-
-    def check_app_installed(self):
-        """Checks if app is installed.
-
-        Must be implemented by subclasses.
-        """
-        raise NotImplementedError()
-
     # Rest of the client methods.
 
-    def start_app(self, wait_time=APP_START_WAIT_TIME):
-        """Starts the server app on the android device.
-
-        Args:
-            wait_time: int, The minimum number of seconds to wait for the app
-                to come up before raising an error. Note that _is_app_running()
-                may take longer than wait_time.
-
-        Raises:
-            AppStartError: When the app was not able to be started.
-        """
-        self.check_app_installed()
-        self._do_start_app()
-        start_time = time.time()
-        expiration_time = start_time + wait_time
-        while time.time() < expiration_time:
-            self._log.debug('Attempting to start %s.', self.app_name)
-            if self._is_app_running():
-                self._log.debug('Successfully started %s after %.1f seconds.',
-                                self.app_name, time.time() - start_time)
-                return
-            time.sleep(1)
-        raise AppStartError('%s failed to start on %s.' % (self.app_name,
-                                                           self._adb.serial))
-
     def connect(self, uid=UNKNOWN_UID, cmd=JsonRpcCommand.INIT):
         """Opens a connection to a JSON RPC server.
 
@@ -232,35 +189,12 @@ class JsonRpcClientBase(object):
         else:
             self.uid = UNKNOWN_UID
 
-    def close(self):
+    def disconnect(self):
         """Close the connection to the remote client."""
         if self._conn:
             self._conn.close()
             self._conn = None
 
-    def _grep(self, regex, output):
-        """Similar to linux's `grep`, this returns the line in an output stream
-        that matches a given regex pattern.
-
-        This function is specifically used to grep strings from AdbProxy's
-        output. We have to do this in Python instead of using cli tools because
-        we need to support windows which does not have `grep` in all vesions.
-
-        Args:
-            regex: string, a regex that matches the expected pattern.
-            output: byte string, the raw output of the adb cmd.
-
-        Returns:
-            A list of strings, all of which are output lines that matches the
-            regex pattern.
-        """
-        lines = output.decode('utf-8').strip().splitlines()
-        results = []
-        for line in lines:
-            if re.search(regex, line):
-                results.append(line.strip())
-        return results
-
     def _cmd(self, command, uid=None):
         """Send a command to the server.
 
@@ -319,20 +253,6 @@ class JsonRpcClientBase(object):
                 method_name=method)
         return result['result']
 
-    def _is_app_running(self):
-        """Checks if the app is currently running on an android device.
-
-        May be overridden by subclasses with custom sanity checks.
-        """
-        running = False
-        try:
-            self.connect()
-            running = True
-        finally:
-            self.close()
-            # This 'return' squashes exceptions from connect()
-            return running
-
     def __getattr__(self, name):
         """Wrapper for python magic to turn method calls into RPC calls."""
 
diff --git a/mobly/controllers/android_device_lib/jsonrpc_shell_base.py b/mobly/controllers/android_device_lib/jsonrpc_shell_base.py
index 1176b4c..e596d8c 100755
--- a/mobly/controllers/android_device_lib/jsonrpc_shell_base.py
+++ b/mobly/controllers/android_device_lib/jsonrpc_shell_base.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 """Shared library for frontends to jsonrpc servers."""
 from __future__ import print_function
 
@@ -23,7 +22,7 @@ from mobly.controllers import android_device
 
 
 class Error(Exception):
-  pass
+    pass
 
 
 class JsonRpcShellBase(object):
@@ -55,7 +54,7 @@ class JsonRpcShellBase(object):
             if len(serials) != 1:
                 raise Error(
                     'Expected one phone, but %d found. Use the -s flag.' %
-                        len(serials))
+                    len(serials))
             serial = serials[0]
         if serial not in serials:
             raise Error('Device "%s" is not found by adb.' % serial)
diff --git a/mobly/controllers/android_device_lib/sl4a_client.py b/mobly/controllers/android_device_lib/sl4a_client.py
index 9e89f09..d961734 100644
--- a/mobly/controllers/android_device_lib/sl4a_client.py
+++ b/mobly/controllers/android_device_lib/sl4a_client.py
@@ -12,51 +12,131 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """JSON RPC interface to android scripting engine."""
+import logging
+import time
 
-import re
-
-from mobly.controllers.android_device_lib import adb
+from mobly import utils
+from mobly.controllers.android_device_lib import event_dispatcher
 from mobly.controllers.android_device_lib import jsonrpc_client_base
 
-DEVICE_SIDE_PORT = 8080
-
+_APP_NAME = 'SL4A'
+_DEVICE_SIDE_PORT = 8080
 _LAUNCH_CMD = (
     'am start -a com.googlecode.android_scripting.action.LAUNCH_SERVER '
     '--ei com.googlecode.android_scripting.extra.USE_SERVICE_PORT %s '
     'com.googlecode.android_scripting/.activity.ScriptingLayerServiceLauncher')
+# Maximum time to wait for the app to start on the device (10 minutes).
+# TODO: This timeout is set high in order to allow for retries in
+# start_app_and_connect. Decrease it when the call to connect() has the option
+# for a quicker timeout than the default _cmd() timeout.
+# TODO: Evaluate whether the high timeout still makes sense for sl4a. It was
+# designed for user snippets which could be very slow to start depending on the
+# size of the snippet and main apps. sl4a can probably use a much smaller value.
+_APP_START_WAIT_TIME = 10 * 60
+
+
+class Error(Exception):
+    pass
+
+
+class AppStartError(Error):
+    """Raised when sl4a is not able to be started."""
 
 
 class Sl4aClient(jsonrpc_client_base.JsonRpcClientBase):
     """A client for interacting with SL4A using Mobly Snippet Lib.
 
-    See superclass documentation for a list of public attributes.
+    Extra public attributes:
+    ed: Event dispatcher instance for this sl4a client.
     """
 
-    def __init__(self, host_port, adb_proxy):
+    def __init__(self, adb_proxy, log=logging.getLogger()):
         """Initializes an Sl4aClient.
 
         Args:
-            host_port: (int) The host port of this RPC client.
-            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
+            self._adb: (adb.AdbProxy) The adb proxy to use to start the app.
+            log: (logging.Logger) logger to which to send log messages.
         """
-        super(Sl4aClient, self).__init__(
-            host_port=host_port,
-            device_port=DEVICE_SIDE_PORT,
-            app_name='SL4A',
-            adb_proxy=adb_proxy)
+        super(Sl4aClient, self).__init__(app_name=_APP_NAME, log=log)
+        self.ed = None
+        self._adb = adb_proxy
 
-    def _do_start_app(self):
+    def start_app_and_connect(self):
         """Overrides superclass."""
+        # Check that sl4a is installed
+        out = self._adb.shell('pm list package')
+        if not utils.grep('com.googlecode.android_scripting', out):
+            raise AppStartError('%s is not installed on %s' %
+                                (_APP_NAME, self._adb.serial))
+
+        # sl4a has problems connecting after disconnection, so kill the apk and
+        # try connecting again.
+        try:
+            self.stop_app()
+        except Exception as e:
+            self.log.warning(e)
+
+        # Launch the app
+        self.host_port = utils.get_available_host_port()
+        self.device_port = _DEVICE_SIDE_PORT
+        self._adb.forward(
+            ['tcp:%d' % self.host_port, 'tcp:%d' % self.device_port])
         self._adb.shell(_LAUNCH_CMD % self.device_port)
 
+        # Connect with retry
+        start_time = time.time()
+        expiration_time = start_time + _APP_START_WAIT_TIME
+        started = False
+        while time.time() < expiration_time:
+            self.log.debug('Attempting to start %s.', self.app_name)
+            try:
+                self.connect()
+                started = True
+                break
+            except:
+                self.log.debug(
+                    '%s is not yet running, retrying',
+                    self.app_name,
+                    exc_info=True)
+            time.sleep(1)
+        if not started:
+            raise jsonrpc_client_base.AppStartError(
+                '%s failed to start on %s.' % (self.app_name, self._adb.serial))
+
+        # Start an EventDispatcher for the current sl4a session
+        event_client = Sl4aClient(self._adb, self.log)
+        event_client.host_port = self.host_port
+        event_client.connect(
+            uid=self.uid, cmd=jsonrpc_client_base.JsonRpcCommand.CONTINUE)
+        self.ed = event_dispatcher.EventDispatcher(event_client)
+        self.ed.start()
+
     def stop_app(self):
         """Overrides superclass."""
-        self._adb.shell('am force-stop com.googlecode.android_scripting')
+        try:
+            if self._conn:
+                # Be polite; let the dest know we're shutting down.
+                try:
+                    self.closeSl4aSession()
+                except:
+                    self.log.exception('Failed to gracefully shut down %s.',
+                                       self.app_name)
 
-    def check_app_installed(self):
-        """Overrides superclass."""
-        out = self._adb.shell('pm list package')
-        if not self._grep('com.googlecode.android_scripting', out):
-            raise jsonrpc_client_base.AppStartError(
-                '%s is not installed on %s' % (self.app_name,
-                                               self._adb.serial))
+                # Close the socket connection.
+                self.disconnect()
+
+                # Close Event Dispatcher
+                if self.ed:
+                    try:
+                        self.ed.clean_up()
+                    except:
+                        self.log.exception(
+                            'Failed to shutdown sl4a event dispatcher.')
+                    self.ed = None
+
+            # Terminate the app
+            self._adb.shell('am force-stop com.googlecode.android_scripting')
+        finally:
+            # Always clean up the adb port
+            if self.host_port:
+                self._adb.forward(['--remove', 'tcp:%d' % self.host_port])
diff --git a/mobly/controllers/android_device_lib/snippet_client.py b/mobly/controllers/android_device_lib/snippet_client.py
index a714cbe..1470927 100644
--- a/mobly/controllers/android_device_lib/snippet_client.py
+++ b/mobly/controllers/android_device_lib/snippet_client.py
@@ -14,6 +14,7 @@
 """JSON RPC interface to Mobly Snippet Lib."""
 import logging
 import re
+import time
 
 from mobly import utils
 from mobly.controllers.android_device_lib import adb
@@ -22,90 +23,144 @@ from mobly.controllers.android_device_lib import jsonrpc_client_base
 _INSTRUMENTATION_RUNNER_PACKAGE = (
     'com.google.android.mobly.snippet.SnippetRunner')
 
-_LAUNCH_CMD = ('am instrument -w -e action start -e port %s %s/' +
-               _INSTRUMENTATION_RUNNER_PACKAGE)
+# 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/' +
+                  _INSTRUMENTATION_RUNNER_PACKAGE)
+
+_LAUNCH_CMD_V1 = (
+    'am instrument -w -e action start %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
 
 _STOP_CMD = (
     'am instrument -w -e action stop %s/' + _INSTRUMENTATION_RUNNER_PACKAGE)
 
+# 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
+
 
 class Error(Exception):
     pass
 
 
+class ProtocolVersionError(Error):
+    """Raised when the protocol reported by the snippet is unknown."""
+
+
 class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
     """A client for interacting with snippet APKs using Mobly Snippet Lib.
 
     See superclass documentation for a list of public attributes.
+
+    It currently supports both v0 and v1 snippet launch protocols, although
+    support for v0 will be removed in a future version.
+
+    For a description of the launch protocols, see the documentation in
+    mobly-snippet-lib, SnippetRunner.java.
     """
 
-    def __init__(self, package, host_port, adb_proxy, log=logging.getLogger()):
+    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.
-            host_port: (int) The port at which to start the snippet client. Note
-                       that the same port will currently be used for both the
-                       device and host side of the connection.
-            adb_proxy: (adb.AdbProxy) The adb proxy to use to start the app.
+            adb_proxy: (adb.AdbProxy) Adb proxy for running adb commands.
+            log: (logging.Logger) logger to which to send log messages.
         """
-        # TODO(adorokhine): Don't assume that a free host-side port is free on
-        # the device as well. Both sides should allocate a unique port.
-        super(SnippetClient, self).__init__(
-            host_port=host_port,
-            device_port=host_port,
-            app_name=package,
-            adb_proxy=adb_proxy,
-            log=log)
+        super(SnippetClient, self).__init__(app_name=package, log=log)
         self.package = package
-        self.log = log
-        self._serial = self._adb.serial
+        self._adb = adb_proxy
         self._proc = None
 
-    def _do_start_app(self):
-        """Overrides superclass."""
-        cmd = _LAUNCH_CMD % (self.device_port, self.package)
-        # 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', self.package)
-        adb_cmd = [adb.ADB]
-        if self._adb.serial:
-            adb_cmd += ['-s', self._adb.serial]
-        adb_cmd += ['shell', cmd]
-        self._proc = utils.start_standing_subprocess(adb_cmd, shell=False)
+    def start_app_and_connect(self):
+        """Overrides superclass. Launches a snippet app and connects to it."""
+        self._check_app_installed()
+
+        # 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
+        start_time = time.time()
+        self._proc = self._do_start_app(cmd)
+
+        # "Instrumentation crashed" could be due to several reasons, eg
+        # exception thrown during startup or just a launch protocol 0 snippet
+        # dying because it needs the port flag. Sadly we have no way to tell so
+        # just warn and retry as v0.
+        # TODO(adorokhine): delete this in Mobly 1.6 when snippet v0 support is
+        # removed.
+        line = self._read_line()
+        if line == 'INSTRUMENTATION_RESULT: shortMsg=Process crashed.':
+            self.log.warning('Snippet %s crashed on startup. This might be an '
+                             'actual error or a snippet using deprecated v0 '
+                             'start protocol. Retrying as a v0 snippet.',
+                             self.package)
+            self.host_port = utils.get_available_host_port()
+            # Reuse the host port as the device port in v0 snippet. This isn't
+            # safe in general, but the protocol is deprecated.
+            cmd = _LAUNCH_CMD_V0 % (self.host_port, self.package)
+            self._proc = self._do_start_app(cmd)
+            self._connect_to_v0()
+        else:
+            # Check protocol version and get the device port
+            match = re.match('^SNIPPET START, PROTOCOL ([0-9]+) ([0-9]+)$',
+                             line)
+            if not match or match.group(1) != '1':
+                raise ProtocolVersionError(line)
+            self._connect_to_v1()
+        self.log.debug('Snippet %s started after %.1fs on host port %s',
+                       self.package, time.time() - start_time, self.host_port)
 
     def stop_app(self):
-        """Overrides superclass."""
         # Kill the pending 'adb shell am instrument -w' process if there is one.
         # Although killing the snippet apk would abort this process anyway, we
         # want to call stop_standing_subprocess() to perform a health check,
         # print the failure stack trace if there was any, and reap it from the
         # process table.
-        if self._proc:
-            utils.stop_standing_subprocess(self._proc)
         self.log.debug('Stopping snippet apk %s', self.package)
-        out = self._adb.shell(_STOP_CMD % self.package).decode('utf-8')
-        if 'OK (0 tests)' not in out:
-            raise Error('Failed to stop existing apk. Unexpected output: %s' %
-                        out)
+        try:
+            # Close the socket connection.
+            self.disconnect()
+            if self._proc:
+                utils.stop_standing_subprocess(self._proc)
+            out = self._adb.shell(_STOP_CMD % self.package).decode('utf-8')
+            if 'OK (0 tests)' not in out:
+                raise Error('Failed to stop existing apk. Unexpected '
+                            'output: %s' % out)
+        finally:
+            # Always clean up the adb port
+            if self.host_port:
+                self._adb.forward(['--remove', 'tcp:%d' % self.host_port])
 
-    def check_app_installed(self):
+    def _start_event_client(self):
         """Overrides superclass."""
+        event_client = SnippetClient(
+            package=self.package,
+            host_port=self.host_port,
+            adb_proxy=self._adb,
+            log=self.log)
+        event_client.connect(self.uid,
+                             jsonrpc_client_base.JsonRpcCommand.CONTINUE)
+        return event_client
+
+    def _check_app_installed(self):
         # Check that the Mobly Snippet app is installed.
         out = self._adb.shell('pm list package')
-        if not self._grep('^package:%s$' % self.package, out):
+        if not utils.grep('^package:%s$' % self.package, out):
             raise jsonrpc_client_base.AppStartError(
-                '%s is not installed on %s' % (self.package, self._serial))
+                '%s is not installed on %s' % (self.package, self._adb.serial))
         # Check that the app is instrumented.
         out = self._adb.shell('pm list instrumentation')
-        matched_out = self._grep('^instrumentation:%s/%s' % (
-            self.package, _INSTRUMENTATION_RUNNER_PACKAGE), out)
+        matched_out = utils.grep('^instrumentation:%s/%s' %
+                                 (self.package,
+                                  _INSTRUMENTATION_RUNNER_PACKAGE), out)
         if not matched_out:
             raise jsonrpc_client_base.AppStartError(
                 '%s is installed on %s, but it is not instrumented.' %
-                (self.package, self._serial))
+                (self.package, self._adb.serial))
         match = re.search('^instrumentation:(.*)\/(.*) \(target=(.*)\)$',
                           matched_out[0])
         target_name = match.group(3)
@@ -113,17 +168,54 @@ class SnippetClient(jsonrpc_client_base.JsonRpcClientBase):
         # same as the snippet package.
         if target_name != self.package:
             out = self._adb.shell('pm list package')
-            if not self._grep('^package:%s$' % target_name, out):
+            if not utils.grep('^package:%s$' % target_name, out):
                 raise jsonrpc_client_base.AppStartError(
                     'Instrumentation target %s is not installed on %s' %
-                    (target_name, self._serial))
+                    (target_name, self._adb.serial))
 
-    def _start_event_client(self):
-        event_client = SnippetClient(
-            package=self.package,
-            host_port=self.host_port,
-            adb_proxy=self._adb,
-            log=self.log)
-        event_client.connect(self.uid,
-                             jsonrpc_client_base.JsonRpcCommand.CONTINUE)
-        return event_client
+    def _do_start_app(self, launch_cmd):
+        adb_cmd = [adb.ADB]
+        if self._adb.serial:
+            adb_cmd += ['-s', self._adb.serial]
+        adb_cmd += ['shell', launch_cmd]
+        return utils.start_standing_subprocess(adb_cmd, shell=False)
+
+    # TODO(adorokhine): delete this in Mobly 1.6 when snippet v0 support is
+    # removed.
+    def _connect_to_v0(self):
+        self.device_port = self.host_port
+        self._adb.forward(
+            ['tcp:%d' % self.host_port, 'tcp:%d' % self.device_port])
+        start_time = time.time()
+        expiration_time = start_time + _APP_START_WAIT_TIME_V0
+        while time.time() < expiration_time:
+            self.log.debug('Attempting to start %s.', self.package)
+            try:
+                self.connect()
+                return
+            except:
+                self.log.debug(
+                    'v0 snippet %s is not yet running, retrying',
+                    self.package,
+                    exc_info=True)
+            time.sleep(1)
+        raise jsonrpc_client_base.AppStartError(
+            '%s failed to start on %s.' % (self.package, self._adb.serial))
+
+    def _connect_to_v1(self):
+        line = self._read_line()
+        match = re.match('^SNIPPET SERVING, PORT ([0-9]+)$', line)
+        if not match:
+            raise ProtocolVersionError(line)
+        self.device_port = int(match.group(1))
+
+        # Forward the device port to a new host port, and connect to that port
+        self.host_port = utils.get_available_host_port()
+        self._adb.forward(
+            ['tcp:%d' % self.host_port, 'tcp:%d' % self.device_port])
+        self.connect()
+
+    def _read_line(self):
+        line = self._proc.stdout.readline().rstrip()
+        self.log.debug('Read line from instrumentation output: "%s"', line)
+        return line
diff --git a/mobly/utils.py b/mobly/utils.py
index 3d01c79..5c998a0 100644
--- a/mobly/utils.py
+++ b/mobly/utils.py
@@ -15,15 +15,14 @@
 import base64
 import concurrent.futures
 import datetime
-import functools
 import logging
 import os
 import platform
 import portpicker
 import psutil
 import random
+import re
 import signal
-import socket
 import string
 import subprocess
 import time
@@ -189,7 +188,7 @@ def find_files(paths, file_predicate):
     file_list = []
     for path in paths:
         p = abs_path(path)
-        for dirPath, subdirList, fileList in os.walk(p):
+        for dirPath, _, fileList in os.walk(p):
             for fname in fileList:
                 name, ext = os.path.splitext(fname)
                 if file_predicate(name, ext):
@@ -243,7 +242,7 @@ def rand_ascii_str(length):
     Returns:
         The random string generated.
     """
-    letters = [random.choice(ascii_letters_and_digits) for i in range(length)]
+    letters = [random.choice(ascii_letters_and_digits) for _ in range(length)]
     return ''.join(letters)
 
 
@@ -421,3 +420,26 @@ def get_available_host_port():
             return port
     raise Error('Failed to find available port after {} retries'.format(
         MAX_PORT_ALLOCATION_RETRY))
+
+
+def grep(regex, output):
+    """Similar to linux's `grep`, this returns the line in an output stream
+    that matches a given regex pattern.
+
+    It does not rely on the `grep` binary and is not sensitive to line endings,
+    so it can be used cross-platform.
+
+    Args:
+        regex: string, a regex that matches the expected pattern.
+        output: byte string, the raw output of the adb cmd.
+
+    Returns:
+        A list of strings, all of which are output lines that matches the
+        regex pattern.
+    """
+    lines = output.decode('utf-8').strip().splitlines()
+    results = []
+    for line in lines:
+        if re.search(regex, line):
+            results.append(line.strip())
+    return results
diff --git a/tools/sl4a_shell.py b/tools/sl4a_shell.py
index 501e656..1006ff8 100755
--- a/tools/sl4a_shell.py
+++ b/tools/sl4a_shell.py
@@ -35,6 +35,7 @@ u'N2F52'
 """
 
 import argparse
+import logging
 
 from mobly.controllers.android_device_lib import jsonrpc_shell_base
 
@@ -42,7 +43,6 @@ from mobly.controllers.android_device_lib import jsonrpc_shell_base
 class Sl4aShell(jsonrpc_shell_base.JsonRpcShellBase):
     def _start_services(self, console_env):
         """Overrides superclass."""
-        self._ad.start_services()
         self._ad.load_sl4a()
         console_env['s'] = self._ad.sl4a
         console_env['sl4a'] = self._ad.sl4a
@@ -64,4 +64,5 @@ if __name__ == '__main__':
         help=
         'Device serial to connect to (if more than one device is connected)')
     args = parser.parse_args()
+    logging.basicConfig(level=logging.INFO)
     Sl4aShell().main(args.serial)


Test patch:
diff --git a/mobly/test_runner.py b/mobly/test_runner.py
index a883a7c..bcdebfa 100644
--- a/mobly/test_runner.py
+++ b/mobly/test_runner.py
@@ -400,11 +400,13 @@ class TestRunner(object):
             raise signals.ControllerError(
                 'Expected to get at least %d controller objects, got %d.' %
                 (min_number, actual_number))
-        self._controller_registry[module_ref_name] = objects
+        # Save a shallow copy of the list for internal usage, so tests can't
+        # affect internal registry by manipulating the object list.
+        self._controller_registry[module_ref_name] = copy.copy(objects)
         # Collect controller information and write to test result.
         # Implementation of 'get_info' is optional for a controller module.
         if hasattr(module, 'get_info'):
-            controller_info = module.get_info(objects)
+            controller_info = module.get_info(copy.copy(objects))
             logging.debug('Controller %s: %s', module_config_name,
                           controller_info)
             self.results.add_controller_info(module_config_name,
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 5adcfc5..3fbe8cf 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
@@ -51,11 +51,7 @@ class MockSocketFile(object):
 
 class FakeRpcClient(jsonrpc_client_base.JsonRpcClientBase):
     def __init__(self):
-        super(FakeRpcClient, self).__init__(
-            host_port=80,
-            device_port=90,
-            app_name='FakeRpcClient',
-            adb_proxy=None)
+        super(FakeRpcClient, self).__init__(app_name='FakeRpcClient')
 
 
 class JsonRpcClientBaseTest(unittest.TestCase):
@@ -164,8 +160,8 @@ class JsonRpcClientBaseTest(unittest.TestCase):
 
         with self.assertRaises(
                 jsonrpc_client_base.ProtocolError,
-                msg=jsonrpc_client_base.ProtocolError.
-                NO_RESPONSE_FROM_HANDSHAKE):
+                msg=
+                jsonrpc_client_base.ProtocolError.NO_RESPONSE_FROM_HANDSHAKE):
             client.some_rpc(1, 2, 3)
 
     @mock.patch('socket.create_connection')
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 21db458..c5b7df4 100755
--- a/tests/mobly/controllers/android_device_lib/snippet_client_test.py
+++ b/tests/mobly/controllers/android_device_lib/snippet_client_test.py
@@ -15,14 +15,11 @@
 from builtins import str
 from builtins import bytes
 
-import json
 import mock
-import socket
 import unittest
 
 from mobly.controllers.android_device_lib import jsonrpc_client_base
 from mobly.controllers.android_device_lib import snippet_client
-from tests.lib import mock_android_device
 
 MOCK_PACKAGE_NAME = 'some.package.name'
 MOCK_MISSING_PACKAGE_NAME = 'not.installed'
@@ -46,17 +43,13 @@ class MockAdbProxy(object):
             if self.apk_not_instrumented:
                 return b''
             if self.target_not_installed:
-                return bytes(
-                    'instrumentation:{p}/{r} (target={mp})'.format(
-                        p=MOCK_PACKAGE_NAME,
-                        r=snippet_client._INSTRUMENTATION_RUNNER_PACKAGE,
-                        mp=MOCK_MISSING_PACKAGE_NAME),
-                    'utf-8')
-            return bytes(
-                'instrumentation:{p}/{r} (target={p})'.format(
+                return bytes('instrumentation:{p}/{r} (target={mp})'.format(
                     p=MOCK_PACKAGE_NAME,
-                    r=snippet_client._INSTRUMENTATION_RUNNER_PACKAGE),
-                'utf-8')
+                    r=snippet_client._INSTRUMENTATION_RUNNER_PACKAGE,
+                    mp=MOCK_MISSING_PACKAGE_NAME), 'utf-8')
+            return bytes('instrumentation:{p}/{r} (target={p})'.format(
+                p=MOCK_PACKAGE_NAME,
+                r=snippet_client._INSTRUMENTATION_RUNNER_PACKAGE), 'utf-8')
 
     def __getattr__(self, name):
         """All calls to the none-existent functions in adb proxy would
@@ -78,42 +71,42 @@ class JsonRpcClientBaseTest(unittest.TestCase):
     @mock.patch(JSONRPC_BASE_PACKAGE)
     def test_check_app_installed_normal(self, mock_create_connection,
                                         mock_client_base):
-        sc = snippet_client.SnippetClient(MOCK_PACKAGE_NAME, 42,
-                                          MockAdbProxy())
-        sc.check_app_installed()
+        sc = self._make_client()
+        sc._check_app_installed()
 
     @mock.patch('socket.create_connection')
     @mock.patch(JSONRPC_BASE_PACKAGE)
     def test_check_app_installed_fail_app_not_installed(
             self, mock_create_connection, mock_client_base):
-        sc = snippet_client.SnippetClient(
-            MOCK_PACKAGE_NAME, 42, MockAdbProxy(apk_not_installed=True))
+        sc = self._make_client(MockAdbProxy(apk_not_installed=True))
         expected_msg = '%s is not installed on .*' % MOCK_PACKAGE_NAME
         with self.assertRaisesRegexp(jsonrpc_client_base.AppStartError,
                                      expected_msg):
-            sc.check_app_installed()
+            sc._check_app_installed()
 
     @mock.patch('socket.create_connection')
     @mock.patch(JSONRPC_BASE_PACKAGE)
     def test_check_app_installed_fail_not_instrumented(
             self, mock_create_connection, mock_client_base):
-        sc = snippet_client.SnippetClient(
-            MOCK_PACKAGE_NAME, 42, MockAdbProxy(apk_not_instrumented=True))
+        sc = self._make_client(MockAdbProxy(apk_not_instrumented=True))
         expected_msg = '%s is installed on .*, but it is not instrumented.' % MOCK_PACKAGE_NAME
         with self.assertRaisesRegexp(jsonrpc_client_base.AppStartError,
                                      expected_msg):
-            sc.check_app_installed()
+            sc._check_app_installed()
 
     @mock.patch('socket.create_connection')
     @mock.patch(JSONRPC_BASE_PACKAGE)
     def test_check_app_installed_fail_target_not_installed(
             self, mock_create_connection, mock_client_base):
-        sc = snippet_client.SnippetClient(
-            MOCK_PACKAGE_NAME, 42, MockAdbProxy(target_not_installed=True))
+        sc = self._make_client(MockAdbProxy(target_not_installed=True))
         expected_msg = 'Instrumentation target %s is not installed on .*' % MOCK_MISSING_PACKAGE_NAME
         with self.assertRaisesRegexp(jsonrpc_client_base.AppStartError,
                                      expected_msg):
-            sc.check_app_installed()
+            sc._check_app_installed()
+
+    def _make_client(self, adb_proxy=MockAdbProxy()):
+        return snippet_client.SnippetClient(
+            package=MOCK_PACKAGE_NAME, adb_proxy=adb_proxy)
 
 
 if __name__ == "__main__":
diff --git a/tests/mobly/test_runner_test.py b/tests/mobly/test_runner_test.py
index 67bf394..843888e 100755
--- a/tests/mobly/test_runner_test.py
+++ b/tests/mobly/test_runner_test.py
@@ -110,6 +110,19 @@ class TestRunnerTest(unittest.TestCase):
         self.assertEqual(magic_devices[0].magic, 'magic1')
         self.assertEqual(magic_devices[1].magic, 'magic2')
 
+    def test_register_controller_change_return_value(self):
+        mock_test_config = self.base_mock_test_config.copy()
+        mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME
+        mock_test_config.controller_configs = {
+            mock_ctrlr_config_name: ['magic1', 'magic2']
+        }
+        tr = test_runner.TestRunner(self.log_dir, self.test_bed_name)
+        magic_devices = tr._register_controller(mock_test_config,
+                                                mock_controller)
+        magic1 = magic_devices.pop(0)
+        self.assertIs(magic1, tr._controller_registry['mock_controller'][0])
+        self.assertEqual(len(tr._controller_registry['mock_controller']), 2)
+
     def test_register_controller_less_than_min_number(self):
         mock_test_config = self.base_mock_test_config.copy()
         mock_ctrlr_config_name = mock_controller.MOBLY_CONTROLLER_CONFIG_NAME