Repository: google/mobly

Issue:
Support marking `UID` for tests
We have long had a `UID` field in our record object.
Let's provide a way for test writer to utilize this field.

Discussion:

Patch:
diff --git a/mobly/records.py b/mobly/records.py
index 1dec134..4048554 100644
--- a/mobly/records.py
+++ b/mobly/records.py
@@ -17,6 +17,7 @@
 import collections
 import copy
 import enum
+import functools
 import io
 import logging
 import sys
@@ -34,6 +35,38 @@ OUTPUT_FILE_DEBUG_LOG = 'test_log.DEBUG'
 OUTPUT_FILE_SUMMARY = 'test_summary.yaml'
 
 
+class Error(Exception):
+    """Raised for errors in record module members."""
+
+
+def uid(uid):
+    """Decorator specifying the unique identifier (UID) of a test case.
+
+    The UID will be recorded in the test's record when executed by Mobly.
+
+    If you use any other decorator for the test method, you may want to use
+    this as the outer-most one.
+
+    Note a common UID system is the Universal Unitque Identifier (UUID), but
+    we are not limiting people to use UUID, hence the more generic name `UID`.
+
+    Args:
+        uid: string, the uid for the decorated test function.
+    """
+    if uid is None:
+        raise ValueError('UID cannot be None.')
+
+    def decorate(test_func):
+        @functools.wraps(test_func)
+        def wrapper(*args, **kwargs):
+            return test_func(*args, **kwargs)
+
+        setattr(wrapper, 'uid', uid)
+        return wrapper
+
+    return decorate
+
+
 class TestSummaryEntryType(enum.Enum):
     """Constants used to identify the type of entries in test summary file.
 


Test patch:
diff --git a/mobly/base_test.py b/mobly/base_test.py
index 88e227b..91a0a4f 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -575,6 +575,7 @@ class BaseTestClass(object):
             test_method: function, The test method to execute.
         """
         tr_record = records.TestResultRecord(test_name, self.TAG)
+        tr_record.uid = getattr(test_method, 'uid', None)
         tr_record.test_begin()
         self.current_test_info = runtime_test_info.RuntimeTestInfo(
             test_name, self.log_path, tr_record)
@@ -664,7 +665,7 @@ class BaseTestClass(object):
         raise Error('"%s" cannot be called outside of %s' %
                     (caller_frames[1][3], expected_func_name))
 
-    def generate_tests(self, test_logic, name_func, arg_sets):
+    def generate_tests(self, test_logic, name_func, arg_sets, uid_func=None):
         """Generates tests in the test class.
 
         This function has to be called inside a test class's
@@ -674,6 +675,9 @@ class BaseTestClass(object):
         parameter sets. This way we reduce code repetition and improve test
         scalability.
 
+        Users can provide an optional function to specify the UID of each test.
+        Not all generated tests are required to have UID.
+
         Args:
             test_logic: function, the common logic shared by all the generated
                 tests.
@@ -682,15 +686,26 @@ class BaseTestClass(object):
                 the test logic function.
             arg_sets: a list of tuples, each tuple is a set of arguments to be
                 passed to the test logic function and name function.
+            uid_func: function, an optional function that takes the same
+                arguments as the test logic function and returns a string that
+                is the corresponding UID.
         """
         self._assert_function_name_in_stack(STAGE_NAME_SETUP_GENERATED_TESTS)
+        root_msg = 'During test generation of "%s":' % test_logic.__name__
         for args in arg_sets:
             test_name = name_func(*args)
             if test_name in self.get_existing_test_names():
                 raise Error(
-                    'Test name "%s" already exists, cannot be duplicated!' %
-                    test_name)
+                    '%s Test name "%s" already exists, cannot be duplicated!' %
+                    (root_msg, test_name))
             test_func = functools.partial(test_logic, *args)
+            if uid_func is not None:
+                uid = uid_func(*args)
+                if uid is None:
+                    logging.warning('%s UID for arg set %s is None.', root_msg,
+                                    args)
+                else:
+                    setattr(test_func, 'uid', uid)
             self._generated_test_table[test_name] = test_func
 
     def _safe_exec_func(self, func, *args):
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index 8d77017..43951f4 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -1855,9 +1855,59 @@ class BaseTestTest(unittest.TestCase):
         bt_cls.run()
         self.assertEqual(len(bt_cls.results.requested), 2)
         self.assertEqual(len(bt_cls.results.passed), 2)
+        self.assertIsNone(bt_cls.results.passed[0].uid)
+        self.assertIsNone(bt_cls.results.passed[1].uid)
         self.assertEqual(bt_cls.results.passed[0].test_name, 'test_1_2')
         self.assertEqual(bt_cls.results.passed[1].test_name, 'test_3_4')
 
+    def test_generate_tests_with_uid(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_generated_tests(self):
+                self.generate_tests(
+                    test_logic=self.logic,
+                    name_func=self.name_gen,
+                    uid_func=self.uid_logic,
+                    arg_sets=[(1, 2), (3, 4)])
+
+            def name_gen(self, a, b):
+                return 'test_%s_%s' % (a, b)
+
+            def uid_logic(self, a, b):
+                return 'uid-%s-%s' % (a, b)
+
+            def logic(self, a, b):
+                pass
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        self.assertEqual(bt_cls.results.passed[0].uid, 'uid-1-2')
+        self.assertEqual(bt_cls.results.passed[1].uid, 'uid-3-4')
+
+    def test_generate_tests_with_none_uid(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            def setup_generated_tests(self):
+                self.generate_tests(
+                    test_logic=self.logic,
+                    name_func=self.name_gen,
+                    uid_func=self.uid_logic,
+                    arg_sets=[(1, 2), (3, 4)])
+
+            def name_gen(self, a, b):
+                return 'test_%s_%s' % (a, b)
+
+            def uid_logic(self, a, b):
+                if a == 1:
+                    return None
+                return 'uid-3-4'
+
+            def logic(self, a, b):
+                pass
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        self.assertIsNone(bt_cls.results.passed[0].uid)
+        self.assertEqual(bt_cls.results.passed[1].uid, 'uid-3-4')
+
     def test_generate_tests_selected_run(self):
         class MockBaseTest(base_test.BaseTestClass):
             def setup_generated_tests(self):
@@ -1925,7 +1975,8 @@ class BaseTestTest(unittest.TestCase):
         self.assertEqual(actual_record.test_name, "setup_generated_tests")
         self.assertEqual(
             actual_record.details,
-            'Test name "ha" already exists, cannot be duplicated!')
+            'During test generation of "logic": Test name "ha" already exists'
+            ', cannot be duplicated!')
         expected_summary = ("Error 1, Executed 0, Failed 0, Passed 0, "
                             "Requested 0, Skipped 0")
         self.assertEqual(bt_cls.results.summary_str(), expected_summary)
@@ -2048,6 +2099,35 @@ class BaseTestTest(unittest.TestCase):
                         'mock_controller: Some failure')
         self.assertEqual(record.details, expected_msg)
 
+    def test_uid(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            @records.uid('some-uid')
+            def test_func(self):
+                pass
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        actual_record = bt_cls.results.passed[0]
+        self.assertEqual(actual_record.uid, 'some-uid')
+
+    def test_uid_not_specified(self):
+        class MockBaseTest(base_test.BaseTestClass):
+            def test_func(self):
+                pass
+
+        bt_cls = MockBaseTest(self.mock_test_cls_configs)
+        bt_cls.run()
+        actual_record = bt_cls.results.passed[0]
+        self.assertIsNone(actual_record.uid)
+
+    def test_uid_is_none(self):
+        with self.assertRaisesRegex(ValueError, 'UID cannot be None.'):
+
+            class MockBaseTest(base_test.BaseTestClass):
+                @records.uid(None)
+                def not_a_test(self):
+                    pass
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/tests/mobly/records_test.py b/tests/mobly/records_test.py
index 9648637..ba25a69 100755
--- a/tests/mobly/records_test.py
+++ b/tests/mobly/records_test.py
@@ -427,6 +427,13 @@ class RecordsTest(unittest.TestCase):
         self.assertEqual(tr.controller_info[0].controller_info,
                          ['magicA', 'magicB'])
 
+    def test_uid(self):
+        @records.uid('some-uuid')
+        def test_uid_helper():
+            """Dummy test used by `test_uid` for testing the uid decorator."""
+
+        self.assertEqual(test_uid_helper.uid, 'some-uuid')
+
 
 if __name__ == '__main__':
     unittest.main()