diff --git a/openhtf/output/callbacks/mfg_inspector.py b/openhtf/output/callbacks/mfg_inspector.py index 896bff0a..b478b598 100644 --- a/openhtf/output/callbacks/mfg_inspector.py +++ b/openhtf/output/callbacks/mfg_inspector.py @@ -14,6 +14,7 @@ """Output and/or upload a TestRun or MfgEvent proto for mfg-inspector.com.""" +import functools import logging import time import zlib @@ -24,6 +25,8 @@ from openhtf.output import callbacks from openhtf.output.proto import test_runs_converter +from google3.googlex.glass.xtf.proto import xtf_pb2 +from openhtf.output.proto import mfg_event_pb2 from openhtf.output.proto import guzzle_pb2 @@ -42,7 +45,7 @@ def _send_mfg_inspector_request( envelope_data: bytes, credentials: credentials_lib.Credentials, destination_url: str, -): +) -> dict[str, ...]: """Send upload http request. Intended to be run in retry loop.""" logging.info('Uploading result...') @@ -73,18 +76,38 @@ def _send_mfg_inspector_request( raise UploadFailedError(message) -def send_mfg_inspector_data(inspector_proto, credentials, destination_url, - payload_type): +@functools.lru_cache(len(guzzle_pb2.PayloadType.values())) +def _is_compressed_payload_type( + payload_type: guzzle_pb2.PayloadType, +) -> bool: + return ( + guzzle_pb2.PayloadType.Name(payload_type) + .lower() + .startswith('compressed_') + ) + + +def send_mfg_inspector_data( + inspector_proto: mfg_event_pb2.MfgEvent | test_runs_pb2.TestRun, + credentials: credentials_lib.Credentials, + destination_url: str, + payload_type: guzzle_pb2.PayloadType, +) -> dict[str, ...]: """Upload MfgEvent to steam_engine.""" envelope = guzzle_pb2.TestRunEnvelope() # pytype: disable=module-attr # gen-stub-imports - envelope.payload = zlib.compress(inspector_proto.SerializeToString()) + data = inspector_proto.SerializeToString() + if _is_compressed_payload_type(payload_type): + data = zlib.compress(data) + + envelope.payload = data envelope.payload_type = payload_type envelope_data = envelope.SerializeToString() for _ in range(5): try: - result = _send_mfg_inspector_request(envelope_data, credentials, - destination_url) + result = _send_mfg_inspector_request( + envelope_data, credentials, destination_url + ) return result except UploadFailedError: time.sleep(1)