Add support for custom units in OpenHTF proto converters.

openhtf/output/proto/mfg_event_converter.py

@@ -409,10 +409,14 @@ def _copy_unidimensional_measurement(self, phase, name, measurement,
     if measurement.docstring:
       mfg_measurement.description = measurement.docstring
     mfg_measurement.parameter_tag.append(phase.name)
-    if (measurement.units and
-        measurement.units.code in test_runs_converter.UOM_CODE_MAP):
-      mfg_measurement.unit_code = (
-          test_runs_converter.UOM_CODE_MAP[measurement.units.code])
+    if measurement.units:
+      if measurement.units.code in test_runs_converter.UOM_CODE_MAP:
+        mfg_measurement.unit_code = test_runs_converter.UOM_CODE_MAP[
+            measurement.units.code
+        ]
+      elif hasattr(mfg_measurement, 'custom_unit_code'):
+        mfg_measurement.custom_unit_code = measurement.units.code or ''
+        mfg_measurement.custom_unit_suffix = measurement.units.suffix or ''
 
     # Copy failed measurements as failure_codes. This happens early to include
     # unset measurements.

openhtf/output/proto/test_runs_converter.py

@@ -208,8 +208,12 @@ def _mangle_measurement(name, measured_value, measurement, mangled_parameters,
     for validator in measurement.validators:
       mangled_param.description += '\nValidator: ' + str(validator)
 
-    if measurement.units and measurement.units.code in UOM_CODE_MAP:
-      mangled_param.unit_code = UOM_CODE_MAP[measurement.units.code]
+    if measurement.units:
+      if measurement.units.code in UOM_CODE_MAP:
+        mangled_param.unit_code = UOM_CODE_MAP[measurement.units.code]
+      elif hasattr(mangled_param, 'custom_unit_code'):
+        mangled_param.custom_unit_code = measurement.units.code or ''
+        mangled_param.custom_unit_suffix = measurement.units.suffix or ''
     mangled_parameters[mangled_name] = mangled_param
 
 
@@ -226,8 +230,12 @@ def _extract_parameters(record, testrun, used_parameter_names):
       testrun_param.name = tr_name
       if measurement.docstring:
         testrun_param.description = measurement.docstring
-      if measurement.units and measurement.units.code in UOM_CODE_MAP:
-        testrun_param.unit_code = UOM_CODE_MAP[measurement.units.code]
+      if measurement.units:
+        if measurement.units.code in UOM_CODE_MAP:
+          testrun_param.unit_code = UOM_CODE_MAP[measurement.units.code]
+        elif hasattr(testrun_param, 'custom_unit_code'):
+          testrun_param.custom_unit_code = measurement.units.code or ''
+          testrun_param.custom_unit_suffix = measurement.units.suffix or ''
 
       if measurement.marginal:
         testrun_param.status = test_runs_pb2.MARGINAL_PASS

test/output/callbacks/callbacks_test.py

@@ -23,14 +23,19 @@
 
 import openhtf as htf
 from openhtf import util
+from openhtf.core import measurements
+from openhtf.core import phase_branches
+from openhtf.core import phase_collections
+from openhtf.core import phase_descriptor
+from openhtf.core import phase_group
+from openhtf.core import test_record as htf_test_record
 from examples import all_the_things
-from openhtf.core import phase_branches, phase_descriptor, phase_collections, phase_group
 from openhtf.output.callbacks import console_summary
 from openhtf.output.callbacks import json_factory
 from openhtf.output.proto import mfg_event_converter
 from openhtf.output.proto import test_runs_converter
-from openhtf.output.proto import test_runs_pb2
 from openhtf.util import test
+from openhtf.output.proto import test_runs_pb2
 
 
 class TestOutput(test.TestCase):
@@ -168,6 +173,48 @@ def test_mfg_event_from_test_record(self, user_mock):
       else:
         raise AssertionError('No attachment named %s' % attachment_name)
 
+  def test_test_run_custom_units(self):
+    custom_unit = util.units.UnitDescriptor(
+        'gibibytes per second', None, 'GiB/s'
+    )
+    measurement = measurements.Measurement('custom-units-meas').with_units(
+        custom_unit
+    )
+    measurement.measured_value.set(10.5)
+    measurement.outcome = measurements.Outcome.PASS
+
+    record = htf_test_record.TestRecord(
+        dut_id='mock-dut-id',
+        station_id='mock-station-id',
+        start_time_millis=100,
+        end_time_millis=500,
+        outcome=htf_test_record.Outcome.PASS,
+    )
+
+    phase = htf_test_record.PhaseRecord(
+        name='mock-phase-name',
+        descriptor_id=1,
+        codeinfo=htf_test_record.CodeInfo.uncaptured(),
+        start_time_millis=200,
+        end_time_millis=400,
+        measurements={
+            'custom-units-meas': measurement,
+        },
+    )
+    record.phases = [phase]
+
+    test_run_proto = test_runs_converter.test_run_from_test_record(record)
+
+    self.assertEqual(len(test_run_proto.test_parameters), 1)
+    param = test_run_proto.test_parameters[0]
+    self.assertEqual(param.name, 'custom-units-meas')
+    self.assertEqual(
+        param.unit_code, test_runs_pb2.Units.UnitCode.Value('NONE')
+    )
+    if hasattr(param, 'custom_unit_code'):
+      self.assertEqual(param.custom_unit_code, '')
+      self.assertEqual(param.custom_unit_suffix, 'GiB/s')
+
 
 class TestConsoleSummary(test.TestCase):
 

test/output/callbacks/mfg_event_converter_test.py

@@ -319,6 +319,34 @@ def test_copy_measurements_from_phase(self):
     self.assertEqual(mock_measurement_within_percent.numeric_marginal_maximum,
                      0)
 
+  def test_copy_measurements_custom_units(self):
+    custom_unit = units.UnitDescriptor('gibibytes per second', None, 'GiB/s')
+    measurement = self._create_and_set_measurement(
+        'custom-units-meas', 10.5
+    ).with_units(custom_unit)
+
+    phase = test_record.PhaseRecord(
+        name='mock-phase-name',
+        descriptor_id=1,
+        codeinfo=self.create_codeinfo(),
+        measurements={
+            'custom-units-meas': measurement,
+        },
+    )
+
+    mfg_event = mfg_event_pb2.MfgEvent()
+    copier = mfg_event_converter.PhaseCopier([phase])
+    copier.copy_measurements(mfg_event)
+
+    self.assertEqual(len(mfg_event.measurement), 1)
+    mfg_measurement = mfg_event.measurement[0]
+    self.assertEqual(mfg_measurement.name, 'custom-units-meas')
+    self.assertEqual(
+        mfg_measurement.unit_code, test_runs_pb2.Units.UnitCode.Value('NONE')
+    )
+    self.assertEqual(mfg_measurement.custom_unit_code, '')
+    self.assertEqual(mfg_measurement.custom_unit_suffix, 'GiB/s')
+
   def test_copy_attachments_from_phase(self):
     first_attachment_name = 'first_attachment_name'
     first_attachment = _create_hacked_massive_attachment()