Fix encoding json with Attachments

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.

Author: Jonathan Van Eenwyk

openhtf/output/callbacks/json_factory.py

@@ -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.
 

openhtf/output/proto/mfg_event_converter.py

@@ -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
@@ -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):

openhtf/output/servers/station_server.py

@@ -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'
 
@@ -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))
 

openhtf/util/json_encoder.py

@@ -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)