From 5f04207463837de325cc69162a5d7e411d802777 Mon Sep 17 00:00:00 2001
From: Dushyant Bhatnagar <dbhatman@google.com>
Date: Fri, 19 Jul 2024 16:26:54 -0700
Subject: [PATCH] Add support for uncompressed uploads if they are being used

PiperOrigin-RevId: 654160330
---
 openhtf/output/callbacks/mfg_inspector.py | 37 +++++++++++++++++++----
 1 file changed, 31 insertions(+), 6 deletions(-)

diff --git a/openhtf/output/callbacks/mfg_inspector.py b/openhtf/output/callbacks/mfg_inspector.py
index 896bff0a..819aa2fb 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,8 +25,12 @@
 from openhtf.output import callbacks
 from openhtf.output.proto import test_runs_converter
 
+from openhtf.output.proto import test_runs_pb2
+from openhtf.output.proto import mfg_event_pb2
 from openhtf.output.proto import guzzle_pb2
 
+from typing import Any, Dict, Union
+
 
 _MFG_INSPECTOR_UPLOAD_TIMEOUT = 60 * 5
 
@@ -42,7 +47,7 @@ def _send_mfg_inspector_request(
     envelope_data: bytes,
     credentials: credentials_lib.Credentials,
     destination_url: str,
-):
+) -> Dict[str, Any]:
   """Send upload http request.  Intended to be run in retry loop."""
   logging.info('Uploading result...')
 
@@ -73,18 +78,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: Union[mfg_event_pb2.MfgEvent, test_runs_pb2.TestRun],
+    credentials: credentials_lib.Credentials,
+    destination_url: str,
+    payload_type: guzzle_pb2.PayloadType,
+) -> Dict[str, Any]:
   """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)