From 8f80a802c5510a556d2ce9d4fc8f7ab9aa2fb7cb Mon Sep 17 00:00:00 2001
From: Reiley Yang <reyang@microsoft.com>
Date: Wed, 15 Aug 2018 13:56:29 -0700
Subject: [PATCH] add tracestate to spancontext

---
 opencensus/trace/span.py              |  8 ++++++++
 opencensus/trace/span_context.py      | 18 ++++++++++--------
 opencensus/trace/trace_options.py     |  7 +++++++
 tests/unit/trace/test_span.py         | 15 +++++++++++++++
 tests/unit/trace/test_span_context.py | 16 ++++++++++------
 5 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/opencensus/trace/span.py b/opencensus/trace/span.py
index 07bb8c212..dcda480ba 100644
--- a/opencensus/trace/span.py
+++ b/opencensus/trace/span.py
@@ -145,6 +145,14 @@ def __init__(
         self._child_spans = []
         self.context_tracer = context_tracer
         self.span_kind = span_kind
+        for callback in Span._on_create_callbacks:
+            callback(self)
+
+    _on_create_callbacks = []
+
+    @staticmethod
+    def on_create(callback):
+        Span._on_create_callbacks.append(callback)
 
     @property
     def children(self):
diff --git a/opencensus/trace/span_context.py b/opencensus/trace/span_context.py
index da412a7b4..a0040949e 100644
--- a/opencensus/trace/span_context.py
+++ b/opencensus/trace/span_context.py
@@ -58,6 +58,7 @@ def __init__(
             trace_id=None,
             span_id=None,
             trace_options=None,
+            tracestate=None,
             from_header=False):
         if trace_id is None:
             trace_id = generate_trace_id()
@@ -69,21 +70,22 @@ def __init__(
         self.trace_id = self._check_trace_id(trace_id)
         self.span_id = self._check_span_id(span_id)
         self.trace_options = trace_options
+        self.tracestate = tracestate
 
-    def __str__(self):
-        """Returns a string form of the SpanContext. This is the format of
-        the Trace Context Header and should be forwarded to downstream
-        requests as the X-Cloud-Trace-Context header.
+    def __repr__(self):
+        """Returns a string form of the SpanContext.
 
         :rtype: str
         :returns: String form of the SpanContext.
         """
-        enabled = self.trace_options.enabled
-        header = '{}/{};o={}'.format(
+        fmt = '{}(trace_id={}, span_id={}, trace_options={}, tracestate={})'
+        return fmt.format(
+            type(self).__name__,
             self.trace_id,
             self.span_id,
-            int(enabled))
-        return header
+            self.trace_options,
+            self.tracestate,
+        )
 
     def _check_span_id(self, span_id):
         """Check the format of the span_id to ensure it is 16-character hex
diff --git a/opencensus/trace/trace_options.py b/opencensus/trace/trace_options.py
index e94adaf0c..a719df630 100644
--- a/opencensus/trace/trace_options.py
+++ b/opencensus/trace/trace_options.py
@@ -44,6 +44,13 @@ def check_trace_options(self, trace_options_byte):
 
         return trace_options_byte
 
+    def __repr__(self):
+        fmt = '{}(enabled={})'
+        return fmt.format(
+            type(self).__name__,
+            self.get_enabled,
+        )
+
     @property
     def get_enabled(self):
         """Get the last bit from the trace options which is the enabled field.
diff --git a/tests/unit/trace/test_span.py b/tests/unit/trace/test_span.py
index 01f147313..868db8739 100644
--- a/tests/unit/trace/test_span.py
+++ b/tests/unit/trace/test_span.py
@@ -212,6 +212,21 @@ def test_finish(self):
         span.finish()
         self.assertIsNotNone(span.end_time)
 
+    def test_on_create(self):
+        from opencensus.trace.span import Span
+        self.on_create_called = False
+        span = self._make_one('span1')
+        self.assertFalse(self.on_create_called)
+        try:
+            @Span.on_create
+            def callback(span):
+                self.on_create_called = True
+            span = self._make_one('span2')
+        finally:
+            Span._on_create_callbacks = []
+        self.assertTrue(self.on_create_called)
+        del self.on_create_called
+
     def test___iter__(self):
         root_span_name = 'root_span_name'
         child1_span_name = 'child1_span_name'
diff --git a/tests/unit/trace/test_span_context.py b/tests/unit/trace/test_span_context.py
index f137081e1..7627c2100 100644
--- a/tests/unit/trace/test_span_context.py
+++ b/tests/unit/trace/test_span_context.py
@@ -15,6 +15,7 @@
 import unittest
 from opencensus.trace import span_context as span_context_module
 from opencensus.trace.trace_options import TraceOptions
+from opencensus.trace.tracestate import Tracestate
 
 
 class TestSpanContext(unittest.TestCase):
@@ -38,17 +39,20 @@ def test_constructor(self):
         self.assertEqual(span_context.trace_id, self.trace_id)
         self.assertEqual(span_context.span_id, self.span_id)
 
-    def test__str__(self):
+    def test__repr__(self):
         span_context = self._make_one(
             trace_id=self.trace_id,
             span_id=self.span_id,
-            trace_options=TraceOptions('1'))
+            trace_options=TraceOptions('1'),
+            tracestate=Tracestate())
 
-        header_expected = '6e0c63257de34c92bf9efcd03927272e' \
-                          '/6e0c63257de34c92;o=1'
-        header = span_context.__str__()
+        expected_repr = 'SpanContext(' \
+                       'trace_id=6e0c63257de34c92bf9efcd03927272e, ' \
+                       'span_id=6e0c63257de34c92, ' \
+                       'trace_options=TraceOptions(enabled=True), ' \
+                       'tracestate=Tracestate())'
 
-        self.assertEqual(header_expected, header)
+        self.assertEqual(expected_repr, span_context.__repr__())
 
     def test_check_span_id_none(self):
         span_context = self._make_one(