diff options
-rw-r--r-- | mobly/controllers/android_device_lib/services/snippet_management_service.py | 19 | ||||
-rwxr-xr-x | tests/mobly/controllers/android_device_lib/services/snippet_management_service_test.py | 51 |
2 files changed, 55 insertions, 15 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 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() |