Squash commit for mfg-inspector output

example/hello_world.py

@@ -12,19 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 """Example OpenHTF test logic.
 
 Run with (your virtualenv must be activated first):
 python ./hello_world.py --config ./hello_world.yaml
 """
 
-
 import tempfile
 import time
 
 import example_plug
 import openhtf
+import openhtf.io.output as output
 
 from openhtf.names import *
 
@@ -39,14 +38,17 @@ def example_monitor(example):
         'widget_type').MatchesRegex(r'.*Widget$').Doc(
             '''This measurement tracks the type of widgets.'''),
     Measurement(
-        'widget_color').Doc('Color of the widget'))
+        'widget_color').Doc('Color of the widget'),
+    Measurement(
+        'widget_size').InRange(0, 10))
 @plug(example=example_plug.Example)
 def hello_world(test, example):
   """A hello world test phase."""
   test.logger.info('Hello World!')
   test.measurements.widget_type = prompts.DisplayPrompt(
       'What\'s the widget type?', text_input=True)
   test.measurements.widget_color = 'Black'
+  test.measurements.widget_size = 5
   test.logger.info('Example says: %s', example.DoStuff())
 
 
@@ -69,7 +71,7 @@ def set_measurements(test):
 @measures(
     Measurement('dimensions').WithDimensions(UOM['HERTZ']),
     Measurement('lots_of_dims').WithDimensions(
-        UOM['HERTZ'], UOM['BYTE'], UOM['RADIAN']))
+        UOM['HERTZ'], UOM['SECOND'], UOM['RADIAN']))
 def dimensions(test):
   for dim in range(5):
     test.measurements.dimensions[dim] = 1 << dim
@@ -85,7 +87,18 @@ def attachments(test):
 
 
 if __name__ == '__main__':
-  test = openhtf.Test(hello_world, set_measurements, dimensions, attachments)
+  test = openhtf.Test(hello_world, set_measurements, dimensions, attachments,
+      # Some metadata fields, these in particular are used by mfg-inspector,
+      # but you can include any metadata fields.
+      test_name='MyTest', test_description='OpenHTF Example Test',
+      test_version='1.0.0')
   test.AddOutputCallback(OutputToJSON(
-  		'./%(dut_id)s.%(start_time_millis)s', indent=4))
+      './%(dut_id)s.%(start_time_millis)s.json', indent=4))
+  test.AddOutputCallback(output.OutputToTestRunProto(
+      './%(dut_id)s.%(start_time_millis)s.json'))
+  # Example of how to upload to mfg-inspector.  Replace user email and key,
+  # these are dummy values.
+  #test.AddOutputCallback(output.UploadToMfgInspector(
+  #  '[email protected]',
+  #  open('my-upload-key.p12', 'r').read()))
   test.Execute(test_start=triggers.PromptForTestStart())

openhtf/__init__.py

@@ -22,7 +22,6 @@
 import signal
 import socket
 import sys
-from json import JSONEncoder
 
 import gflags
 import mutablerecords
@@ -31,7 +30,6 @@
 from openhtf import exe
 from openhtf import plugs
 from openhtf import util
-from openhtf.exe import test_state
 from openhtf.exe import triggers
 from openhtf.io import http_api
 from openhtf.io import rundata
@@ -47,43 +45,6 @@ class InvalidTestPhaseError(Exception):
   """Raised when an invalid method is decorated."""
 
 
-class OutputToJSON(JSONEncoder):
-  """Return an output callback that writes JSON Test Records.
-
-  An example filename_pattern might be:
-    '/data/test_records/%(dut_id)s.%(start_time_millis)s'
-
-  To use this output mechanism:
-    test = openhtf.Test(PhaseOne, PhaseTwo)
-    test.AddOutputCallback(openhtf.OutputToJson(
-        '/data/test_records/%(dut_id)s.%(start_time_millis)s'))
-
-  Args:
-    filename_pattern: A format string specifying the filename to write to,
-      will be formatted with the Test Record as a dictionary.
-  """
-
-  def __init__(self, filename_pattern, **kwargs):
-    super(OutputToJSON, self).__init__(**kwargs)
-    self.filename_pattern = filename_pattern
-
-  def default(self, obj):
-    # Handle a few custom objects that end up in our output.
-    if isinstance(obj, BaseException):
-      # Just repr exceptions.
-      return repr(obj)
-    if isinstance(obj, conf.Config):
-      return obj.dictionary
-    if obj in test_state.TestState.State:
-      return str(obj)
-    return super(OutputToJSON, self).default(obj)
-
-  def __call__(self, test_record):  # pylint: disable=invalid-name
-    as_dict = util.convert_to_dict(test_record)
-    with open(self.filename_pattern % as_dict, 'w') as f:  # pylint: disable=invalid-name
-      f.write(self.encode(as_dict))
-
-
 class TestPhaseOptions(mutablerecords.Record(
     'TestPhaseOptions', [], {'timeout_s': None, 'run_if': None})):
   """
@@ -144,14 +105,16 @@ class Test(object):
 
   Args:
     *phases: The ordered list of phases to execute for this test.
+    **metadata: Any metadata that should be associated with test records.
   """
 
-  def __init__(self, *phases):
+  def __init__(self, *phases, **metadata):
     """Creates a new Test to be executed.
 
     Args:
       *phases: The ordered list of phases to execute for this test.
     """
+    self.metadata = metadata
     self.loop = False
     self.phases = [TestPhaseInfo.WrapOrReturn(phase) for phase in phases]
     self.output_callbacks = []
@@ -182,6 +145,7 @@ def AddOutputCallback(self, callback):
 
   def OutputTestRecord(self, test_record):
     """Feed the record of this test to all output modules."""
+    test_record.metadata.update(self.metadata)
     for output_cb in self.output_callbacks:
       output_cb(test_record)
 

openhtf/exe/phase_data.py

@@ -34,6 +34,7 @@
 
 _LOG = logging.getLogger(__name__)
 
+
 class DuplicateAttachmentError(Exception):
   """Raised when two attachments are attached with the same name."""
 
@@ -125,6 +126,7 @@ def RecordPhaseTiming(self, phase, test_state):
         measurement.name: copy.deepcopy(measurement)
         for measurement in phase.measurements
     }
+
     # Populate dummy declaration list for frontend API.
     test_state.running_phase.measurements = {
         measurement.name: measurement._asdict()

openhtf/exe/test_state.py

@@ -126,14 +126,13 @@ def SetStateRunning(self):
 
   def SetStateFinished(self):
     """Mark the state as finished, only called if the test ended normally."""
-    if any(meas.outcome == 'FAIL'
+    if any(not meas.outcome
            for phase in self.record.phases
            for meas in phase.measurements.itervalues()):
       self._state = self.State.FAIL
     else:
       self._state = self.State.PASS
 
-
   def GetFinishedRecord(self):
     """Get a test_record.TestRecord for the finished test.
 
@@ -157,7 +156,7 @@ def GetFinishedRecord(self):
           'Blank or missing DUT ID, HTF requires a non-blank ID.')
 
     self.record.end_time_millis = util.TimeMillis()
-    self.record.outcome = self._state
+    self.record.outcome = str(self._state)
     return self.record
 
   def __str__(self):

openhtf/io/output/__init__.py

@@ -0,0 +1,90 @@
+"""Output module for outputting to JSON."""
+
+import logging
+import oauth2client.client
+import threading
+
+from openhtf import util
+from openhtf.io.output import json_factory
+from openhtf.io.output import mfg_inspector
+
+
+class OutputToTestRunProto(object):  # pylint: disable=too-few-public-methods
+  """Return an output callback that writes mfg-inspector TestRun Protos.
+
+  An example filename_pattern might be:
+    '/data/test_records/%(dut_id)s.%(start_time_millis)s'
+
+  To use this output mechanism:
+    test = openhtf.Test(PhaseOne, PhaseTwo)
+    test.AddOutputCallback(openhtf.OutputToTestRunProto(
+        '/data/test_records/%(dut_id)s.%(start_time_millis)s'))
+
+  Args:
+    filename_pattern: A format string specifying the filename to write to,
+      will be formatted with the Test Record as a dictionary.
+  """
+
+  def __init__(self, filename_pattern):
+    self.filename_pattern = filename_pattern
+
+  def __call__(self, test_record):  # pylint: disable=invalid-name
+    as_dict = util.convert_to_dict(test_record)
+    with open(self.filename_pattern % as_dict, 'w') as outfile:
+      outfile.write(mfg_inspector.TestRunFromTestRecord(
+          test_record).SerializeToString())
+
+
+class UploadToMfgInspector(object):  # pylint: disable=too-few-public-methods
+  """Generate a mfg-inspector TestRun proto and upload it.
+
+  Create an output callback to upload to mfg-inspector.com using the given
+  username and authentication key (which should be the key data itself, not a
+  filename or file).
+  """
+
+  TOKEN_URI = 'https://accounts.google.com/o/oauth2/token'
+  SCOPE_CODE_URI = 'https://www.googleapis.com/auth/glass.infra.quantum_upload'
+  DESTINATION_URL = ('https://clients2.google.com/factoryfactory/'
+                     'uploads/quantum_upload/')
+
+  # pylint: disable=invalid-name,missing-docstring
+  class _MemStorage(oauth2client.client.Storage):
+    """Helper Storage class that keeps credentials in memory."""
+    def __init__(self):
+      self._lock = threading.Lock()
+      self._credentials = None
+
+    def acquire_lock(self):
+      self._lock.acquire(True)
+
+    def release_lock(self):
+      self._lock.release()
+
+    def locked_get(self):
+      return self._credentials
+
+    def locked_put(self, credentials):
+      self._credentials = credentials
+  # pylint: enable=invalid-name,missing-docstring
+
+  def __init__(self, user, keydata):
+    self.user = user
+    self.keydata = keydata
+
+  def __call__(self, test_record):  # pylint: disable=invalid-name
+    credentials = oauth2client.client.SignedJwtAssertionCredentials(
+        service_account_name=self.user,
+        private_key=self.keydata,
+        scope=self.SCOPE_CODE_URI,
+        user_agent='OpenHTF Guzzle Upload Client',
+        token_uri=self.TOKEN_URI)
+    credentials.set_store(self._MemStorage())
+
+    testrun = mfg_inspector.TestRunFromTestRecord(test_record)
+    try:
+      mfg_inspector.UploadTestRun(testrun, self.DESTINATION_URL, credentials)
+    except mfg_inspector.UploadFailedError:
+      # For now, just log the exception.  Once output is a bit more robust,
+      # we can propagate this up and handle it accordingly.
+      logging.exception('Upload to mfg-inspector failed!')

openhtf/io/output/json_factory.py

@@ -0,0 +1,53 @@
+"""Module for outputting test record to JSON-formatted files."""
+
+from json import JSONEncoder
+
+from openhtf import conf
+from openhtf import util
+from openhtf.exe import test_state
+
+
+class OutputToJSON(JSONEncoder):
+  """Return an output callback that writes JSON Test Records.
+
+  An example filename_pattern might be:
+    '/data/test_records/%(dut_id)s.%(start_time_millis)s'
+
+  To use this output mechanism:
+    test = openhtf.Test(PhaseOne, PhaseTwo)
+    test.AddOutputCallback(openhtf.OutputToJson(
+        '/data/test_records/%(dut_id)s.%(start_time_millis)s'))
+
+  Args:
+    filename_pattern: A format string specifying the filename to write to,
+      will be formatted with the Test Record as a dictionary.
+    inline_attachments: Whether attachments should be included inline in the
+      output.  Set to False if you expect to have large binary attachments.
+  """
+
+  def __init__(self, filename_pattern=None, inline_attachments=True, **kwargs):
+    super(OutputToJSON, self).__init__(**kwargs)
+    self.filename_pattern = filename_pattern
+    self.inline_attachments = inline_attachments
+
+  def default(self, obj):
+    # Handle a few custom objects that end up in our output.
+    if isinstance(obj, BaseException):
+      # Just repr exceptions.
+      return repr(obj)
+    if isinstance(obj, conf.Config):
+      return obj.dictionary
+    if obj in test_state.TestState.State:
+      return str(obj)
+    return super(OutputToJSON, self).default(obj)
+
+  # pylint: disable=invalid-name
+  def __call__(self, test_record):
+    assert self.filename_pattern, 'filename_pattern required'
+    if self.inline_attachments:
+      as_dict = util.convert_to_dict(test_record)
+    else:
+      as_dict = util.convert_to_dict(test_record, ignore_keys='attachments')
+    with open(self.filename_pattern % as_dict, 'w') as f:
+      f.write(self.encode(as_dict))
+  # pylint: enable=invalid-name