Add more test cases

Author: ocelotl

lightstep/b3_propagator.py

@@ -4,8 +4,8 @@
 
 from basictracer.propagator import Propagator
 from basictracer.context import SpanContext
-from opentracing import span
 from opentracing import SpanContext as OTSpanContext
+from opentracing import SpanContextCorruptedException
 
 _LOG = getLogger(__name__)
 _SINGLE_HEADER = "b3"
@@ -37,7 +37,7 @@ def inject(self, span_context, carrier):
 
         flags = baggage.pop(_FLAGS, None)
         if flags is not None:
-            carrier[self._FLAGS] = flags
+            carrier[_FLAGS] = flags
 
         sampled = baggage.pop(_SAMPLED, None)
         if sampled is not None:
@@ -56,9 +56,7 @@ def inject(self, span_context, carrier):
                     )
                 carrier[_SAMPLED] = int(sampled)
 
-        if (sampled is None and flags is None) and (
-            traceid is None or spanid is None
-        ):
+        if sampled is flags is (traceid and spanid) is None:
             warn(
                 "If not propagating only the sampling state, traceid and "
                 "spanid must be defined"
@@ -92,30 +90,36 @@ def extract(self, carrier):
         baggage = {}
 
         if _SINGLE_HEADER in carrier.keys():
-            fields = carrier[_SINGLE_HEADER].split("-", 4)
+            fields = carrier.pop(_SINGLE_HEADER).split("-", 4)
+            baggage.update(carrier)
             len_fields = len(fields)
-            if len(fields) == 1:
+            if len_fields == 1:
                 sampled = fields[0]
-            elif len(fields) == 2:
+            elif len_fields == 2:
                 traceid, spanid = fields
-            elif len(fields) == 3:
+            elif len_fields == 3:
                 traceid, spanid, sampled = fields
             else:
                 traceid, spanid, sampled, parent_spanid = fields
-                baggage[_PARENTSPANID] = parent_spanid
+                baggage[_PARENTSPANID] = int(parent_spanid, 16)
             if sampled == "d":
                 baggage[_FLAGS] = 1
             else:
-                baggage[_SAMPLED] = sampled
+                baggage[_SAMPLED] = int(sampled, 16)
         else:
             traceid = carrier.pop(_TRACEID, None)
             spanid = carrier.pop(_SPANID, None)
             parentspanid = carrier.pop(_PARENTSPANID, None)
             sampled = carrier.pop(_SAMPLED, None)
             flags = carrier.pop(_FLAGS, None)
 
+            if sampled is flags is (traceid and spanid) is None:
+
+                raise SpanContextCorruptedException()
+
             if parentspanid is not None:
-                baggage[_PARENTSPANID] = parentspanid
+                baggage[_PARENTSPANID] = int(parentspanid, 16)
+
             if flags == 1:
                 baggage[_FLAGS] = flags
                 if sampled is not None:
@@ -124,8 +128,10 @@ def extract(self, carrier):
                         "the received value of x-b3-sampled"
                     )
             elif sampled is not None:
-                baggage[_SAMPLED] = sampled
+                baggage[_SAMPLED] = int(sampled, 16)
+
             baggage.update(carrier)
+
             if baggage == OTSpanContext.EMPTY_BAGGAGE:
                 baggage = None
         return SpanContext(

tests/b3_propagator_test.py

@@ -1,45 +1,125 @@
-import unittest
+from unittest import TestCase
+
+from pytest import raises
+from opentracing import SpanContextCorruptedException
 
-import lightstep
 from lightstep.propagation import B3
+from lightstep import Tracer
 
 
-class B3PropagatorTest(unittest.TestCase):
+class B3PropagatorTest(TestCase):
     def setUp(self):
-        self._tracer = lightstep.Tracer(
-                periodic_flush_seconds=0,
-                collector_host='localhost')
+        self._tracer = Tracer(
+            periodic_flush_seconds=0,
+            collector_host="localhost"
+        )
 
     def tracer(self):
         return self._tracer
 
     def tearDown(self):
         self._tracer.flush()
 
-    def testInjectExtract(self):
+    def test_inject(self):
         carrier = {}
-        span = self.tracer().start_span('Sending request')
-
-        span.set_baggage_item('checked', 'baggage')
+        span = self.tracer().start_span("test_inject")
+        span.set_baggage_item("checked", "baggage")
+        self.tracer().inject(span.context, B3, carrier)
+        self.assertEqual(
+            carrier, {
+                "x-b3-traceid": format(span.context.trace_id, "032x"),
+                "x-b3-spanid": format(span.context.span_id, "016x"),
+                "checked": "baggage"
+            }
+        )
 
+        carrier = {}
+        span = self.tracer().start_span("test_inject")
+        span.set_baggage_item("x-b3-flags", 1)
+        span.set_baggage_item("x-b3-sampled", 0)
         self.tracer().inject(span.context, B3, carrier)
+        self.assertEqual(
+            carrier, {
+                "x-b3-traceid": format(span.context.trace_id, "032x"),
+                "x-b3-spanid": format(span.context.span_id, "016x"),
+                "x-b3-flags": 1,
+            }
+        )
+
+    def test_extract_multiple_headers(self):
+
+        result = self.tracer().extract(
+            B3, {
+                "x-b3-traceid": format(12, "032x"),
+                "x-b3-spanid": format(345, "016x"),
+                "checked": "baggage"
+            }
+        )
+
+        self.assertEqual(12, result.trace_id)
+        self.assertEqual(345, result.span_id)
+        self.assertEqual({"checked": "baggage"}, result.baggage)
+
+        result = self.tracer().extract(
+            B3, {
+                "x-b3-traceid": format(12, "032x"),
+                "x-b3-spanid": format(345, "016x"),
+                "x-b3-flags": 1,
+                "x-b3-sampled": 0
+            }
+        )
+
+        self.assertEqual(12, result.trace_id)
+        self.assertEqual(345, result.span_id)
+        self.assertEqual({"x-b3-flags": 1}, result.baggage)
 
-        result = self.tracer().extract(B3, carrier)
-        self.assertEqual(span.context.span_id, result.span_id)
-        self.assertEqual(span.context.trace_id, result.trace_id)
-        self.assertEqual(span.context.baggage, result.baggage)
-        self.assertEqual(span.context.sampled, result.sampled)
+    def test_extract_single_header(self):
+        result = self.tracer().extract(
+            B3, {
+                "b3": "a12-b34-1-c56",
+                "checked": "baggage"
+            }
+        )
+        self.assertEqual(2578, result.trace_id)
+        self.assertEqual(2868, result.span_id)
+        self.assertDictEqual(
+            {
+                "x-b3-sampled": 1,
+                "x-b3-parentspanid": 3158,
+                "checked": "baggage"
+            }, result.baggage
+        )
 
-    def single_header_extraction(self):
-        # FIXME complete this test case
-        input_ = 1
+        result = self.tracer().extract(
+            B3, {
+                "b3": "a12-b34-d-c56",
+                "checked": "baggage"
+            }
+        )
+        self.assertEqual(2578, result.trace_id)
+        self.assertEqual(2868, result.span_id)
+        self.assertDictEqual(
+            {
+                "x-b3-flags": 1,
+                "x-b3-parentspanid": 3158,
+                "checked": "baggage"
+            }, result.baggage
+        )
 
-        result = self.tracer().extract(B3, bytearray(input, 'utf-8'))
-        self.assertEqual(6397081719746291766, result.span_id)
-        self.assertEqual(506100417967962170, result.trace_id)
-        self.assertEqual(True, result.sampled)
-        self.assertDictEqual(result.baggage, {"checked" : "baggage"})
+    def test_invalid_traceid_spanid(self):
 
+        with raises(SpanContextCorruptedException):
+            self.tracer().extract(
+                B3, {
+                    "x-b3-spanid": format(345, "016x"),
+                    "checked": "baggage"
+                }
+            )
 
-if __name__ == '__main__':
-    unittest.main()
+        with raises(SpanContextCorruptedException):
+            self.tracer().extract(
+                B3, {
+                    "x-b3-traceid": format(345, "032x"),
+                    "checked": "baggage"
+                }
+            )