Skip to content

Commit

Permalink
Add support for uncompressed uploads if they are being used
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 654160330
  • Loading branch information
dbhatman authored and copybara-github committed Jul 22, 2024
1 parent 1f103eb commit 9225cbb
Showing 1 changed file with 29 additions and 6 deletions.
35 changes: 29 additions & 6 deletions openhtf/output/callbacks/mfg_inspector.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand All @@ -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...')

Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit 9225cbb

Please sign in to comment.