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

1 files changed, 38 insertions, 13 deletions

diff --git a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py b/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
index 54bf60d..27d389a 100755
--- a/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
+++ b/tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py
@@ -162,23 +162,41 @@ class SnippetManagementServiceTest(unittest.TestCase):
     manager.foo.ha('param')
     mock_client.ha.assert_called_once_with('param')
 
-  def test_client_v2_flag_default_value(self):
-    mock_device = mock.MagicMock()
-    mock_device.dimensions = {}
-    manager = snippet_management_service.SnippetManagementService(mock_device)
-    self.assertFalse(manager._is_using_client_v2())
+  # TODO(mhaoli): The client v2 switch is transient, we will remove related
+  # tests after we complete the migration from v1 to v2.
+  def test_client_v2_switch_default_value(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=False)
+
+  def test_client_v2_switch_when_set_device_dimension_to_false(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=False,
+        dimensions={'use_mobly_snippet_client_v2': 'false'})
 
-  def test_client_v2_flag_false(self):
-    mock_device = mock.MagicMock(
+  def test_client_v2_switch_when_set_device_dimension_to_true(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=True,
+        dimensions={'use_mobly_snippet_client_v2': 'true'})
+
+  def test_client_v2_switch_when_set_device_attribute_to_false(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=False, use_mobly_snippet_client_v2='false')
+
+  def test_client_v2_switch_when_set_device_attribute_to_true(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=True, use_mobly_snippet_client_v2='true')
+
+  def test_client_v2_switch_when_both_attribute_and_dimension_are_false(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=False,
+        use_mobly_snippet_client_v2='false',
         dimensions={'use_mobly_snippet_client_v2': 'false'})
-    manager = snippet_management_service.SnippetManagementService(mock_device)
-    self.assertFalse(manager._is_using_client_v2())
 
-  def test_client_v2_flag_true(self):
-    mock_device = mock.MagicMock(
+  def test_client_v2_switch_when_both_attribute_and_dimension_are_true(self):
+    self._set_device_attribute_and_check_client_v2_switch(
+        expect_switch_value=True,
+        use_mobly_snippet_client_v2='true',
         dimensions={'use_mobly_snippet_client_v2': 'true'})
-    manager = snippet_management_service.SnippetManagementService(mock_device)
-    self.assertTrue(manager._is_using_client_v2())
 
   @mock.patch(SNIPPET_CLIENT_V2_CLASS_PATH)
   def test_client_v2_add_snippet_client(self, mock_class):
@@ -261,6 +279,13 @@ class SnippetManagementServiceTest(unittest.TestCase):
 
     mock_client.restore_server_connection.assert_called_once_with()
 
+  def _set_device_attribute_and_check_client_v2_switch(self,
+                                                       expect_switch_value,
+                                                       **device_attributes):
+    mock_device = mock.MagicMock(**device_attributes)
+    manager = snippet_management_service.SnippetManagementService(mock_device)
+    self.assertEqual(expect_switch_value, manager._is_using_client_v2())
+
 
 if __name__ == '__main__':
   unittest.main()