diff --git a/mobly/base_test.py b/mobly/base_test.py
index 00ba3d9..1742f19 100644
--- a/mobly/base_test.py
+++ b/mobly/base_test.py
@@ -838,21 +838,22 @@
         self.current_test_info = None
     return tr_record
 
-  def _assert_function_name_in_stack(self, expected_func_name):
-    """Asserts that the current stack contains the given function name."""
+  def _assert_function_names_in_stack(self, expected_func_names):
+    """Asserts that the current stack contains any of the given function names.
+    """
     current_frame = inspect.currentframe()
     caller_frames = inspect.getouterframes(current_frame, 2)
     for caller_frame in caller_frames[2:]:
-      if caller_frame[3] == expected_func_name:
+      if caller_frame[3] in expected_func_names:
         return
-    raise Error('"%s" cannot be called outside of %s' %
-                (caller_frames[1][3], expected_func_name))
+    raise Error(f"'{caller_frames[1][3]}' cannot be called outside of the "
+                f"following functions: {expected_func_names}.")
 
   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
-    `self.setup_generated_tests` function.
+    This function has to be called inside a test class's `self.pre_run` or
+    `self.setup_generated_tests`.
 
     Generated tests are not written down as methods, but as a list of
     parameter sets. This way we reduce code repetition and improve test
@@ -873,7 +874,8 @@
         arguments as the test logic function and returns a string that
         is the corresponding UID.
     """
-    self._assert_function_name_in_stack(STAGE_NAME_PRE_RUN)
+    self._assert_function_names_in_stack(
+        [STAGE_NAME_PRE_RUN, 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)
@@ -926,7 +928,7 @@
     'test_*'.
 
     Note this only gets the names of tests that already exist. If
-    `setup_generated_test` has not happened when this was called, the
+    `generate_tests` has not happened when this was called, the
     generated tests won't be listed.
 
     Returns:
diff --git a/tests/mobly/base_test_test.py b/tests/mobly/base_test_test.py
index 5cc6b20..5754dd7 100755
--- a/tests/mobly/base_test_test.py
+++ b/tests/mobly/base_test_test.py
@@ -1991,6 +1991,31 @@
     self.assertEqual(bt_cls.results.skipped, [])
 
   # TODO(angli): remove after the full deprecation of `setup_generated_tests`.
+  def test_setup_generated_tests(self):
+
+    class MockBaseTest(base_test.BaseTestClass):
+
+      def setup_generated_tests(self):
+        self.generate_tests(test_logic=self.logic,
+                            name_func=self.name_gen,
+                            arg_sets=[(1, 2), (3, 4)])
+
+      def name_gen(self, a, b):
+        return 'test_%s_%s' % (a, b)
+
+      def logic(self, a, b):
+        pass
+
+    bt_cls = MockBaseTest(self.mock_test_cls_configs)
+    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')
+
+  # TODO(angli): remove after the full deprecation of `setup_generated_tests`.
   def test_setup_generated_tests_failure(self):
     """Test code path for setup_generated_tests failure.
 
@@ -2136,8 +2161,10 @@
     actual_record = bt_cls.results.error[0]
     utils.validate_test_result(bt_cls.results)
     self.assertEqual(actual_record.test_name, "test_ha")
-    self.assertEqual(actual_record.details,
-                     '"generate_tests" cannot be called outside of pre_run')
+    self.assertEqual(
+        actual_record.details,
+        "'generate_tests' cannot be called outside of the followin"
+        "g functions: ['pre_run', 'setup_generated_tests'].")
     expected_summary = ("Error 1, Executed 1, Failed 0, Passed 0, "
                         "Requested 1, Skipped 0")
     self.assertEqual(bt_cls.results.summary_str(), expected_summary)