diff --git a/openhtf/output/proto/mfg_event_converter.py b/openhtf/output/proto/mfg_event_converter.py index af987a480..a1be181bf 100644 --- a/openhtf/output/proto/mfg_event_converter.py +++ b/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. diff --git a/openhtf/output/proto/test_runs_converter.py b/openhtf/output/proto/test_runs_converter.py index f849525f5..ee9843eaa 100644 --- a/openhtf/output/proto/test_runs_converter.py +++ b/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 diff --git a/test/output/callbacks/callbacks_test.py b/test/output/callbacks/callbacks_test.py index faa22326f..6d387504b 100644 --- a/test/output/callbacks/callbacks_test.py +++ b/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): diff --git a/test/output/callbacks/mfg_event_converter_test.py b/test/output/callbacks/mfg_event_converter_test.py index 1e0be8663..bc3b8505a 100644 --- a/test/output/callbacks/mfg_event_converter_test.py +++ b/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()