1 files changed, 53 insertions, 0 deletions
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index c8f5a1e..ebac0e0 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -2291,6 +2291,27 @@ class BaseTestTest(unittest.TestCase):
     for i, record in enumerate(bt_cls.results.passed):
       self.assertEqual(record.test_name, f'test_something_{i}')
 
+  def test_repeat_generated(self):
+    repeat_count = 3
+
+    class MockBaseTest(base_test.BaseTestClass):
+
+      @base_test.repeat(count=repeat_count)
+      def _run_test_logic(self, arg):
+        pass
+
+      def setup_generated_tests(self):
+        self.generate_tests(
+          self._run_test_logic,
+          name_func=lambda arg: f'test_generated_{arg}',
+          arg_sets=[(1,)])
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run()
+    self.assertEqual(repeat_count, len(bt_cls.results.passed))
+    for i, record in enumerate(bt_cls.results.passed):
+      self.assertEqual(record.test_name, f'test_generated_1_{i}')
+
   def test_repeat_with_failures(self):
     repeat_count = 3
     mock_action = mock.MagicMock()
@@ -2456,6 +2477,38 @@ class BaseTestTest(unittest.TestCase):
     self.assertIs(error_record_1, error_record_2.retry_parent)
     self.assertIs(error_record_2, pass_record.retry_parent)
 
+  def test_retry_generated_test_last_pass(self):
+    max_count = 3
+    mock_action = mock.MagicMock(
+      side_effect = [Exception('Fail 1'), Exception('Fail 2'), None])
+
+    class MockBaseTest(base_test.BaseTestClass):
+
+      @base_test.retry(max_count=max_count)
+      def _run_test_logic(self, arg):
+        mock_action()
+
+      def setup_generated_tests(self):
+        self.generate_tests(
+          self._run_test_logic,
+          name_func=lambda arg: f'test_generated_{arg}',
+          arg_sets=[(1,)])
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    bt_cls.run()
+    self.assertTrue(bt_cls.results.is_all_pass,
+                    'This test run should be considered pass.')
+    self.assertEqual(3, len(bt_cls.results.executed))
+    self.assertEqual(1, len(bt_cls.results.passed))
+    pass_record = bt_cls.results.passed[0]
+    self.assertEqual(pass_record.test_name, f'test_generated_1_retry_2')
+    self.assertEqual(2, len(bt_cls.results.error))
+    error_record_1, error_record_2 = bt_cls.results.error
+    self.assertEqual(error_record_1.test_name, 'test_generated_1')
+    self.assertEqual(error_record_2.test_name, 'test_generated_1_retry_1')
+    self.assertIs(error_record_1, error_record_2.retry_parent)
+    self.assertIs(error_record_2, pass_record.retry_parent)
+
   def test_retry_all_fail(self):
     max_count = 3
     mock_action = mock.MagicMock()