Add BaseSpan interface to make Span and BlankSpan compatible (#313)

mayurkale22 · web-flow · commit cac4305e9c1e · 2018-09-18T22:06:46.000-07:00
* Add base span interface to make Span and BlankSpan compatible

* Add test for context_manager and minor import style fix
diff --git a/opencensus/trace/base_span.py b/opencensus/trace/base_span.py
@@ -0,0 +1,101 @@
+# Copyright 2018, OpenCensus Authors
+#
+# 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.
+
+"""Module containing base class for Span."""
+
+
+class BaseSpan(object):
+    """Base class for Opencensus spans.
+    Subclasses of :class:`BaseSpan` must implement the below methods.
+    """
+
+    @staticmethod
+    def on_create(callback):
+        raise NotImplementedError
+
+    @property
+    def children(self):
+        """The child spans of the current span."""
+        raise NotImplementedError
+
+    def span(self, name='child_span'):
+        """Create a child span for the current span and append it to the child
+        spans list.
+
+        :type name: str
+        :param name: (Optional) The name of the child span.
+
+        :rtype: :class: `~opencensus.trace.span.Span`
+        :returns: A child Span to be added to the current span.
+        """
+        raise NotImplementedError
+
+    def add_attribute(self, attribute_key, attribute_value):
+        """Add attribute to span.
+
+        :type attribute_key: str
+        :param attribute_key: Attribute key.
+
+        :type attribute_value:str
+        :param attribute_value: Attribute value.
+        """
+        raise NotImplementedError
+
+    def add_annotation(self, description, **attrs):
+        """Add an annotation to span.
+
+        :type description: str
+        :param description: A user-supplied message describing the event.
+                        The maximum length for the description is 256 bytes.
+
+        :type attrs: kwargs
+        :param attrs: keyworded arguments e.g. failed=True, name='Caching'
+        """
+        raise NotImplementedError
+
+    def add_time_event(self, time_event):
+        """Add a TimeEvent.
+
+        :type time_event: :class: `~opencensus.trace.time_event.TimeEvent`
+        :param time_event: A TimeEvent object.
+        """
+        raise NotImplementedError
+
+    def add_link(self, link):
+        """Add a Link.
+
+        :type link: :class: `~opencensus.trace.link.Link`
+        :param link: A Link object.
+        """
+        raise NotImplementedError
+
+    def start(self):
+        """Set the start time for a span."""
+        raise NotImplementedError
+
+    def finish(self):
+        """Set the end time for a span."""
+        raise NotImplementedError
+
+    def __iter__(self):
+        """Iterate through the span tree."""
+        raise NotImplementedError
+
+    def __enter__(self):
+        """Start a span."""
+        raise NotImplementedError
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        """Finish a span."""
+        raise NotImplementedError
diff --git a/opencensus/trace/blank_span.py b/opencensus/trace/blank_span.py
@@ -12,12 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from opencensus.trace import base_span
 from opencensus.trace import time_event as time_event_module
 from opencensus.trace.span_context import generate_span_id
 from opencensus.trace.tracers import base
 
 
-class BlankSpan(object):
+class BlankSpan(base_span.BaseSpan):
     """A BlankSpan is an individual timed event which forms a node of the trace
     tree. All operations are no-op.
 
@@ -72,6 +73,10 @@ def __init__(
         self.context_tracer = context_tracer
         self.span_kind = span_kind
 
+    @staticmethod
+    def on_create(callback):
+        pass
+
     @property
     def children(self):
         """The child spans of the current BlankSpan."""
@@ -143,3 +148,11 @@ def finish(self):
     def __iter__(self):
         """Iterate through the span tree."""
         yield self
+
+    def __enter__(self):
+        """Start a span."""
+        return self
+
+    def __exit__(self, exception_type, exception_value, traceback):
+        """Finish a span."""
+        pass
diff --git a/opencensus/trace/span.py b/opencensus/trace/span.py
@@ -16,6 +16,7 @@
 from itertools import chain
 
 from opencensus.trace import attributes
+from opencensus.trace import base_span
 from opencensus.trace import link as link_module
 from opencensus.trace import stack_trace
 from opencensus.trace import status
@@ -31,7 +32,7 @@ class SpanKind(object):
     CLIENT = 2
 
 
-class Span(object):
+class Span(base_span.BaseSpan):
     """A span is an individual timed event which forms a node of the trace
     tree. Each span has its name, span id and parent id. The parent id
     indicates the causal relationships between the individual spans in a
diff --git a/tests/unit/trace/test_base_span.py b/tests/unit/trace/test_base_span.py
@@ -0,0 +1,101 @@
+# Copyright 2018, OpenCensus Authors
+#
+# 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 unittest
+import mock
+from opencensus.trace.base_span import BaseSpan
+
+
+class TestBaseSpan(unittest.TestCase):
+
+    def test_span_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.span('root_span')
+
+    def test_children_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.children
+
+    def test_create_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            @BaseSpan.on_create
+            def callback(span):
+                pass
+
+    def test_start_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.start()
+
+    def test_finish_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.finish()
+
+    def test_add_attribute_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.add_attribute("key", "value")
+
+    def test_add_annotation_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.add_annotation("desc")
+
+    def test_add_time_event_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.add_time_event(None)
+
+    def test_add_link_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            span.add_link(None)
+
+    def test_iter_abstract(self):
+        span = BaseSpan()
+
+        with self.assertRaises(NotImplementedError):
+            list(iter(span))
+
+    @mock.patch.object(BaseSpan, '__exit__')
+    @mock.patch.object(BaseSpan, '__enter__')
+    def test_context_manager_called(self, mock_enter, mock_exit):
+        span = BaseSpan()
+        with span:
+            pass
+        self.assertTrue(mock_enter.called)
+        self.assertTrue(mock_exit.called)
+
+    def test_context_manager_methods(self):
+        span = BaseSpan()
+        with self.assertRaises(NotImplementedError):
+            span.__enter__()
+
+        with self.assertRaises(NotImplementedError):
+            span.__exit__(None, None, None)
+
+
diff --git a/tests/unit/trace/test_blank_span.py b/tests/unit/trace/test_blank_span.py
@@ -1,8 +1,20 @@
-import unittest
-
-import mock
+# Copyright 2018, OpenCensus Authors
+#
+# 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 datetime
+import mock
+import unittest
 from opencensus.trace.link import Link
 from opencensus.trace.span import format_span_json
 from opencensus.trace.time_event import TimeEvent
@@ -71,3 +83,90 @@ def test_do_not_crash(self):
 
         span.start()
         span.finish()
+
+    def test_constructor_explicit(self):
+        from datetime import datetime
+
+        span_id = 'test_span_id'
+        span_name = 'test_span_name'
+        parent_span = mock.Mock()
+        start_time = datetime.utcnow().isoformat() + 'Z'
+        end_time = datetime.utcnow().isoformat() + 'Z'
+        attributes = {
+            'http.status_code': '200',
+            'component': 'HTTP load balancer',
+        }
+        time_events = mock.Mock()
+        links = mock.Mock()
+        stack_trace = mock.Mock()
+        status = mock.Mock()
+        context_tracer = mock.Mock()
+
+        span = self._make_one(
+            name=span_name,
+            parent_span=parent_span,
+            attributes=attributes,
+            start_time=start_time,
+            end_time=end_time,
+            span_id=span_id,
+            stack_trace=stack_trace,
+            time_events=time_events,
+            links=links,
+            status=status,
+            context_tracer=context_tracer)
+
+        self.assertEqual(span.name, span_name)
+        self.assertIsNotNone(span.span_id)
+        self.assertEqual(span.attributes, {})
+        self.assertEqual(span.start_time, start_time)
+        self.assertEqual(span.end_time, end_time)
+        self.assertEqual(span.time_events, time_events)
+        self.assertEqual(span.stack_trace, stack_trace)
+        self.assertEqual(span.links, [])
+        self.assertEqual(span.status, status)
+        self.assertEqual(span.children, [])
+        self.assertEqual(span.context_tracer, context_tracer)
+
+    def test_start(self):
+        span_name = 'root_span'
+        span = self._make_one(span_name)
+        self.assertIsNone(span.start_time)
+
+        span.start()
+        self.assertIsNone(span.start_time)
+
+    def test_finish_without_context_tracer(self):
+        span_name = 'root_span'
+        span = self._make_one(span_name)
+        self.assertIsNone(span.end_time)
+
+        span.finish()
+        self.assertIsNone(span.end_time)
+
+    def test_finish(self):
+        span_name = 'root_span'
+        span = self._make_one(span_name)
+        self.assertIsNone(span.end_time)
+
+        span.finish()
+        self.assertIsNone(span.end_time)
+
+    def test_on_create(self):
+        from opencensus.trace.blank_span import BlankSpan
+        self.on_create_called = False
+        span = self._make_one('span1')
+        self.assertFalse(self.on_create_called)
+        try:
+            @BlankSpan.on_create
+            def callback(span):
+                self.on_create_called = True
+            span = self._make_one('span2')
+        finally:
+            BlankSpan._on_create_callbacks = []
+        self.assertFalse(self.on_create_called)
+
+    def test_context_manager(self):
+        span_name = 'root_span'
+        with self._make_one(span_name) as s:
+            self.assertIsNotNone(s)
+            self.assertEquals(s.name, span_name)
