Have the assertMeasured methods accept the desired outcome as an optional argument, making them more suitable for parameterized tests.

PiperOrigin-RevId: 853837558

Author: glados-verma

openhtf/util/test.py

@@ -925,7 +925,14 @@ def _check_phase(phase_record, strict=False):
         _check_phase(phase_record)
 
   @_assert_phase_or_test_record
-  def assertMeasured(self, phase_record, measurement, value=mock.ANY):
+  def assertMeasured(
+      self,
+      phase_record,
+      measurement,
+      value=mock.ANY,
+      outcome: measurements.Outcome | None = None,
+  ):
+    """Asserts that the given measurement is set; value and outcome optional."""
     self.assertIn(measurement, phase_record.measurements,
                   f'Measurement {measurement} not found')
     self.assertTrue(
@@ -937,12 +944,22 @@ def assertMeasured(self, phase_record, measurement, value=mock.ANY):
           'Measurement %s has wrong value: expected %s, got %s' %
           (measurement, value,
            phase_record.measurements[measurement].measured_value.value))
+    if outcome is not None:
+      self.assertIs(
+          outcome,
+          phase_record.measurements[measurement].outcome,
+      )
 
   @_assert_phase_or_test_record
   def assertMeasuredAlmostEqual(
-      self, phase_record, measurement, value, delta=None
+      self,
+      phase_record,
+      measurement,
+      value,
+      delta=None,
+      outcome: measurements.Outcome | None = None,
   ):
-    self.assertMeasured(phase_record, measurement)
+    self.assertMeasured(phase_record, measurement, mock.ANY, outcome)
     measured_value = phase_record.measurements[measurement].measured_value.value
     self.assertAlmostEqual(
         value,
@@ -956,15 +973,15 @@ def assertMeasuredAlmostEqual(
 
   @_assert_phase_or_test_record
   def assertMeasurementPass(self, phase_record, measurement, value=mock.ANY):
-    self.assertMeasured(phase_record, measurement, value)
-    self.assertIs(measurements.Outcome.PASS,
-                  phase_record.measurements[measurement].outcome)
+    self.assertMeasured(
+        phase_record, measurement, value, measurements.Outcome.PASS
+    )
 
   @_assert_phase_or_test_record
   def assertMeasurementFail(self, phase_record, measurement, value=mock.ANY):
-    self.assertMeasured(phase_record, measurement, value)
-    self.assertIs(measurements.Outcome.FAIL,
-                  phase_record.measurements[measurement].outcome)
+    self.assertMeasured(
+        phase_record, measurement, value, measurements.Outcome.FAIL
+    )
 
   @_assert_phase_or_test_record
   def assertMeasurementMarginal(

test/util/test_test.py

@@ -22,9 +22,9 @@
 from unittest import mock
 
 import openhtf
-
 from openhtf import plugs
 from openhtf.core import base_plugs
+from openhtf.core import measurements
 from openhtf.util import test
 from openhtf.util import validators
 
@@ -162,7 +162,17 @@ def _run_my_phase_in_test_asserts(self, mock_my_plug, test_record):
     self.assertTestOutcomeCode(test_record, 0xBED)
     self.assertNotMeasured(test_record, 'unset_measurement')
     self.assertNotMeasured(test_record.phases[-1], 'unset_measurement')
-    self.assertMeasured(test_record, 'test_measurement', _DO_STUFF_RETVAL)
+    with self.subTest('assertMeasured no value'):
+      self.assertMeasured(test_record, 'test_measurement')
+    with self.subTest('assertMeasured with value'):
+      self.assertMeasured(test_record, 'test_measurement', _DO_STUFF_RETVAL)
+    with self.subTest('assertMeasured with value and outcome'):
+      self.assertMeasured(
+          test_record,
+          'test_measurement',
+          _DO_STUFF_RETVAL,
+          measurements.Outcome.PASS,
+      )
     self.assertMeasured(test_record, 'othr_measurement', 0xDEAD)
     with self.subTest(name='assert_measurement_pass_without_value'):
       self.assertMeasurementPass(test_record, 'passes')
@@ -174,10 +184,27 @@ def _run_my_phase_in_test_asserts(self, mock_my_plug, test_record):
       self.assertMeasurementFail(test_record, 'fails', 20)
     with self.subTest(name='assert_measurement_almost_equal'):
       self.assertMeasuredAlmostEqual(test_record, 'numeric_measurement', 10.0)
+    with self.subTest(name='assert_measurement_almost_equal_with_outcome'):
+      self.assertMeasuredAlmostEqual(
+          test_record,
+          'numeric_measurement',
+          10.0,
+          measurements.Outcome.PASS,
+      )
     with self.subTest(name='assert_measurement_almost_equal_with_delta'):
       self.assertMeasuredAlmostEqual(
           test_record, 'numeric_measurement', 9.5, delta=1.0
       )
+    with self.subTest(
+        name='assert_measurement_almost_equal_with_delta_and_outcome'
+    ):
+      self.assertMeasuredAlmostEqual(
+          test_record,
+          'numeric_measurement',
+          9.5,
+          delta=1.0,
+          outcome=measurements.Outcome.PASS,
+      )
 
   def test_execute_phase_or_test_test_with_patched_plugs(self):
     self.auto_mock_plugs(MyPlug)