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 25, 2020
1 parent 95f15c4 commit 367c467
Show file tree
Hide file tree
Showing 4 changed files with 37 additions and 11 deletions.
11 changes: 1 addition & 10 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.json_encoder import TestRecordEncoder
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
3 changes: 2 additions & 1 deletion 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.json_encoder import TestRecordEncoder


from past.builtins import unicode
Expand Down Expand Up @@ -172,7 +173,7 @@ 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)
json_encoder = TestRecordEncoder(sort_keys=True, indent=2, ensure_ascii=False)
pieces = []
for piece in json_encoder.iterencode(obj):
if isinstance(piece, bytes):
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.json_encoder import TestRecordEncoder

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 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
28 changes: 28 additions & 0 deletions openhtf/util/json_encoder.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
# 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 367c467

Please sign in to comment.