[telemetry] Add ability to add metadata to the trace data

telemetry/telemetry/timeline/trace_data_unittest.py

Index: telemetry/telemetry/timeline/trace_data_unittest.py
diff --git a/telemetry/telemetry/timeline/trace_data_unittest.py b/telemetry/telemetry/timeline/trace_data_unittest.py
index d776bd27724ec0d7b73b88933b795c3bc173e88f..b5c631d7f5ba242d8543853d15a450bddf059779 100644
--- a/telemetry/telemetry/timeline/trace_data_unittest.py
+++ b/telemetry/telemetry/timeline/trace_data_unittest.py
@@ -3,6 +3,8 @@
 # found in the LICENSE file.
 
 import cStringIO
+import datetime
+import exceptions
 import json
 import os
 import tempfile
@@ -98,3 +100,43 @@ class TraceDataBuilderTest(unittest.TestCase):
     self.assertTrue(d.HasTraceFor(trace_data.BATTOR_TRACE_PART))
 
     self.assertRaises(Exception, builder.AsData)
+
+  def testSetTraceFor(self):
+    telemetry_trace = {
+        'traceEvents': [1, 2, 3],
+        'metadata': {
+          'field1': 'value1'
+        }
+    }
+
+    builder = trace_data.TraceDataBuilder()
+    builder.SetTraceFor(trace_data.TELEMETRY_PART, telemetry_trace)
+    d = builder.AsData()
+
+    self.assertEqual(d.GetTraceFor(trace_data.TELEMETRY_PART), telemetry_trace)
+
+  def testSetTraceForRaisesWithInvalidPart(self):
+    builder = trace_data.TraceDataBuilder()
+
+    self.assertRaises(exceptions.AssertionError,
+                      lambda: builder.SetTraceFor('not_a_trace_part', {}))
+
+  def testSetTraceForRaisesWithInvalidTrace(self):
+    builder = trace_data.TraceDataBuilder()
+
+    self.assertRaises(exceptions.AssertionError, lambda:
+        builder.SetTraceFor(trace_data.TELEMETRY_PART, datetime.time.min))
+
+  def testSetTraceForRaisesWithAlreadySetPart(self):
+    builder = trace_data.TraceDataBuilder()
+    builder.SetTraceFor(trace_data.TELEMETRY_PART, {})
+
+    self.assertRaises(exceptions.Exception,
+        lambda: builder.SetTraceFor(trace_data.TELEMETRY_PART, {}))
+
+  def testSetTraceForRaisesAfterAsData(self):
+    builder = trace_data.TraceDataBuilder()
+    builder.AsData()
+
+    self.assertRaises(exceptions.Exception,
+        lambda: builder.SetTraceFor(trace_data.TELEMETRY_PART, {}))