Skip to content

Commit

Permalink
Fix encoding json with Attachments
Browse files Browse the repository at this point in the history
When using Attachments, we need to properly encode them into json to
share them outside. OutputToJSON handled this correctly, but both
MfgInspector and StationServer did not. Move TestRecordEncoder into
openhtf.util.json, so that it can be shared by all three.  For
station_server, patch the tornado json encoder, so that all values being
sent to the client through tornado are properly encoded.
  • Loading branch information
Jonathan Van Eenwyk committed Aug 27, 2020
1 parent 95f15c4 commit 846fc8d
Show file tree
Hide file tree
Showing 4 changed files with 40 additions and 13 deletions.
13 changes: 2 additions & 11 deletions openhtf/output/callbacks/json_factory.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,10 @@
from openhtf.core import test_record
from openhtf.output import callbacks
from openhtf.util import data
from openhtf.util import json_encoder
import six


class TestRecordEncoder(json.JSONEncoder):

def default(self, obj):
if isinstance(obj, test_record.Attachment):
dct = obj._asdict()
dct['data'] = base64.standard_b64encode(obj.data).decode('utf-8')
return dct
return super(TestRecordEncoder, self).default(obj)


class OutputToJSON(callbacks.OutputToFile):
"""Return an output callback that writes JSON Test Records.
Expand Down Expand Up @@ -47,7 +38,7 @@ def __init__(self, filename_pattern=None, inline_attachments=True, **kwargs):
# Conform strictly to the JSON spec by default.
kwargs.setdefault('allow_nan', False)
self.allow_nan = kwargs['allow_nan']
self.json_encoder = TestRecordEncoder(**kwargs)
self.json_encoder = json_encoder.TestRecordEncoder(**kwargs)

def serialize_test_record(self, test_record):
return self.json_encoder.iterencode(self.convert_to_dict(test_record))
Expand Down
5 changes: 3 additions & 2 deletions openhtf/output/proto/mfg_event_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from openhtf.util import data as htf_data
from openhtf.util import units
from openhtf.util import validators
from openhtf.util import json_encoder


from past.builtins import unicode
Expand Down Expand Up @@ -172,9 +173,9 @@ def _convert_object_to_json(obj):
# Since there will be parts of this that may have unicode, either as
# measurement or in the logs, we have to be careful and convert everything
# to unicode, merge, then encode to UTF-8 to put it into the proto.
json_encoder = json.JSONEncoder(sort_keys=True, indent=2, ensure_ascii=False)
encoder = json_encoder.TestRecordEncoder(sort_keys=True, indent=2, ensure_ascii=False)
pieces = []
for piece in json_encoder.iterencode(obj):
for piece in encoder.iterencode(obj):
if isinstance(piece, bytes):
pieces.append(unicode(piece, errors='replace'))
else:
Expand Down
6 changes: 6 additions & 0 deletions openhtf/output/servers/station_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from openhtf.util import logs
from openhtf.util import multicast
from openhtf.util import timeouts
from openhtf.util import json_encoder

STATION_SERVER_TYPE = 'station'

Expand Down Expand Up @@ -556,6 +557,11 @@ def __init__(self, history_path=None):
if not tornado_logger.handlers:
tornado_logger.addHandler(logging.NullHandler())

# Override tornado's json encoding to handle our Attachments.
def _json_encode(value):
return json_encoder.TestRecordEncoder().encode(value)
sockjs.tornado.proto.json_encode = _json_encode

# Bind port early so that the correct port number can be used in the routes.
sockets, port = web_gui_server.bind_port(int(conf.station_server_port))

Expand Down
29 changes: 29 additions & 0 deletions openhtf/util/json_encoder.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
# Copyright 2016 Google Inc. All Rights Reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import base64
import json

from openhtf.core import test_record


class TestRecordEncoder(json.JSONEncoder):
"""JSON encoder that supports Attachments."""

def default(self, obj):
if isinstance(obj, test_record.Attachment):
dct = obj._asdict()
dct['data'] = base64.standard_b64encode(obj.data).decode('utf-8')
return dct
return super(TestRecordEncoder, self).default(obj)

0 comments on commit 846fc8d

Please sign in to comment.