Add assertion that verifies and yields the desired phase record

We often use this pattern when writing unit tests against a test record and need to find a phase given its name.

PiperOrigin-RevId: 743225123

Author: glados-verma

openhtf/util/test.py

@@ -127,6 +127,7 @@ def test_multiple(self, mock_my_plug):
 """
 
 from collections.abc import Callable as CollectionsCallable, Iterator
+import contextlib
 import functools
 import inspect
 import logging
@@ -786,6 +787,21 @@ def assertTestOutcomeCode(self, test_rec, code):
         any(details.code == code for details in test_rec.outcome_details),
         'No OutcomeDetails had code %s' % code)
 
+  @contextlib.contextmanager
+  def assertTestHasPhaseRecord(self, test_rec, phase_name):
+    """Yields a PhaseRecord with the given name, else asserts."""
+    all_phase_names = []
+    expected_phase_rec = None
+    for phase_rec in test_rec.phases:
+      all_phase_names.append(phase_rec.name)
+      if phase_rec.name == phase_name:
+        expected_phase_rec = phase_rec
+    self.assertIsNotNone(
+        expected_phase_rec,
+        msg=f'Phase "{phase_name}" not found in test phases: {all_phase_names}',
+    )
+    yield expected_phase_rec
+
   ##### PhaseRecord Assertions #####
 
   def assertPhaseContinue(self, phase_record):

test/util/test_test.py

@@ -278,6 +278,19 @@ def bad_test(cls_self):  # pylint: disable=unused-argument
     with self.assertRaises(test.InvalidTestError):
       test.yields_phases(bad_test)(self)
 
+  def test_assert_test_has_phase_record(self):
+    self.auto_mock_plugs(MyPlug)
+    test_record = self.execute_phase_or_test(
+        openhtf.Test(test_phase_with_shameless_plug, test_phase)
+    )
+    with self.subTest('phase_found'):
+      with self.assertTestHasPhaseRecord(test_record, test_phase.name) as phase:
+        self.assertEqual(phase.name, test_phase.name)
+    with self.subTest('phase_not_found'):
+      with self.assertRaises(AssertionError):
+        with self.assertTestHasPhaseRecord(test_record, 'nonexistent_phase'):
+          pass
+
 
 class PhaseProfilingTest(test.TestCase):
   """Test profiling an OpenHTF phase in unit testing.