Support switching to client V2 via Android Device property setting (#835)

1 files changed, 17 insertions, 2 deletions

diff --git a/mobly/controllers/android_device_lib/services/snippet_management_service.py b/mobly/controllers/android_device_lib/services/snippet_management_service.py
index 7d1dd94..9dc9d75 100644
--- a/mobly/controllers/android_device_lib/services/snippet_management_service.py
+++ b/mobly/controllers/android_device_lib/services/snippet_management_service.py
@@ -19,6 +19,10 @@ from mobly.controllers.android_device_lib.services import base_service
 
 MISSING_SNIPPET_CLIENT_MSG = 'No snippet client is registered with name "%s".'
 
+# This config is transient and we will remove it after completing the migration
+# from v1 to v2.
+_CLIENT_V2_CONFIG_KEY = 'use_mobly_snippet_client_v2'
+
 
 class Error(errors.ServiceError):
   """Root error type for snippet management service."""
@@ -40,6 +44,8 @@ class SnippetManagementService(base_service.BaseService):
     super().__init__(device)
     self._use_client_v2_switch = None
 
+  # TODO(mhaoli): The client v2 switch is transient, we will remove it after we
+  # complete the migration from v1 to v2.
   def _is_using_client_v2(self):
     """Is this service using snippet client V2.
 
@@ -49,11 +55,20 @@ class SnippetManagementService(base_service.BaseService):
 
     NOTE: This is a transient function when we are migrating the snippet client
     from v1 to v2. It will be removed after the migration is completed.
+
+    Returns:
+      A bool for whether this service is using snippet client V2.
     """
     if self._use_client_v2_switch is None:
       device_dimensions = getattr(self._device, 'dimensions', {})
-      self._use_client_v2_switch = (device_dimensions.get(
-          'use_mobly_snippet_client_v2', 'false').lower() == 'true')
+      switch_from_dimension = (device_dimensions.get(_CLIENT_V2_CONFIG_KEY,
+                                                     'false').lower() == 'true')
+
+      switch_from_attribute = (getattr(self._device, _CLIENT_V2_CONFIG_KEY,
+                                       'false').lower() == 'true')
+
+      self._use_client_v2_switch = (switch_from_dimension or
+                                    switch_from_attribute)
     return self._use_client_v2_switch
 
   @property