Repository: census-instrumentation/opencensus-python
Pull request #169 (2018-04-30 21:27:13)

Add streaming support and message events to gRPC interceptors
The callback based util being used before didn't work very well for streams. Refactoring based off of the example in https://github.com/grpc/grpc/pull/14306.

wkiser: @liyanhui1228, any idea whats wrong with the most recent circleci checks?
liyanhui1228: @wkiser System tests are flaky sometimes, rebuilding it works.

Patch:
diff --git a/opencensus/trace/ext/grpc/client_interceptor.py b/opencensus/trace/ext/grpc/client_interceptor.py
index 1691b10..c8160fd 100644
--- a/opencensus/trace/ext/grpc/client_interceptor.py
+++ b/opencensus/trace/ext/grpc/client_interceptor.py
@@ -14,12 +14,15 @@
 
 import collections
 import logging
+
 import grpc
 import six
 
-from opencensus.trace import execution_context
 from opencensus.trace import attributes_helper
+from opencensus.trace import execution_context
+from opencensus.trace import time_event
 from opencensus.trace.ext import grpc as oc_grpc
+from opencensus.trace.ext.grpc import utils as grpc_utils
 from opencensus.trace.propagation import binary_format
 
 log = logging.getLogger(__name__)
@@ -91,7 +94,8 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
         self.tracer.end_span()
 
     def _intercept_call(
-            self, client_call_details, request_iterator, grpc_type):
+        self, client_call_details, request_iterator, grpc_type
+    ):
         metadata = ()
         if client_call_details.metadata is not None:
             metadata = client_call_details.metadata
@@ -118,10 +122,21 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
             metadata,
             client_call_details.credentials)
 
+        request_iterator = grpc_utils.wrap_iter_with_message_events(
+            request_or_response_iter=request_iterator,
+            span=current_span,
+            message_event_type=time_event.Type.SENT
+        )
+
         return client_call_details, request_iterator, current_span
 
     def _callback(self, current_span):
         def callback(future_response):
+            grpc_utils.add_message_event(
+                proto_message=future_response.result(),
+                span=current_span,
+                message_event_type=time_event.Type.RECEIVED,
+            )
             execution_context.set_current_span(current_span)
             self._trace_future_exception(future_response)
             self.tracer.end_span()
@@ -141,7 +156,8 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
             attribute_value=exception)
 
     def intercept_unary_unary(
-            self, continuation, client_call_details, request):
+        self, continuation, client_call_details, request
+    ):
         if CLOUD_TRACE in client_call_details.method:
             response = continuation(client_call_details, request)
             return response
@@ -159,8 +175,9 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
         return response
 
-    def intercept_unary_stream(self, continuation, client_call_details,
-                               request):
+    def intercept_unary_stream(
+        self, continuation, client_call_details, request
+    ):
         if CLOUD_TRACE in client_call_details.method:
             response = continuation(client_call_details, request)
             return response
@@ -173,12 +190,18 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
         response_it = continuation(
             new_details,
             next(new_request_iterator))
-        self._end_span_between_context(current_span)
+        response_it = grpc_utils.wrap_iter_with_message_events(
+            request_or_response_iter=response_it,
+            span=current_span,
+            message_event_type=time_event.Type.RECEIVED
+        )
+        response_it = grpc_utils.wrap_iter_with_end_span(response_it)
 
         return response_it
 
-    def intercept_stream_unary(self, continuation, client_call_details,
-                               request_iterator):
+    def intercept_stream_unary(
+        self, continuation, client_call_details, request_iterator
+    ):
         if CLOUD_TRACE in client_call_details.method:
             response = continuation(client_call_details, request_iterator)
             return response
@@ -196,8 +219,9 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
 
         return response
 
-    def intercept_stream_stream(self, continuation, client_call_details,
-                                request_iterator):
+    def intercept_stream_stream(
+        self, continuation, client_call_details, request_iterator
+    ):
         if CLOUD_TRACE in client_call_details.method:
             response = continuation(client_call_details, request_iterator)
             return response
@@ -210,7 +234,12 @@ class OpenCensusClientInterceptor(grpc.UnaryUnaryClientInterceptor,
         response_it = continuation(
             new_details,
             new_request_iterator)
-        self._end_span_between_context(current_span)
+        response_it = grpc_utils.wrap_iter_with_message_events(
+            request_or_response_iter=response_it,
+            span=current_span,
+            message_event_type=time_event.Type.RECEIVED
+        )
+        response_it = grpc_utils.wrap_iter_with_end_span(response_it)
 
         return response_it
 
diff --git a/opencensus/trace/ext/grpc/server_interceptor.py b/opencensus/trace/ext/grpc/server_interceptor.py
index 275dcc7..56bbf03 100644
--- a/opencensus/trace/ext/grpc/server_interceptor.py
+++ b/opencensus/trace/ext/grpc/server_interceptor.py
@@ -11,19 +11,20 @@
 # 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 collections
 import logging
 import sys
 
-from google.rpc import code_pb2
 import grpc
+from google.rpc import code_pb2
 
 from opencensus.trace import attributes_helper
 from opencensus.trace import execution_context
 from opencensus.trace import stack_trace as stack_trace
 from opencensus.trace import status
+from opencensus.trace import time_event
 from opencensus.trace import tracer as tracer_module
 from opencensus.trace.ext import grpc as oc_grpc
+from opencensus.trace.ext.grpc import utils as grpc_utils
 from opencensus.trace.propagation import binary_format
 
 ATTRIBUTE_COMPONENT = 'COMPONENT'
@@ -31,76 +32,66 @@ ATTRIBUTE_ERROR_NAME = 'ERROR_NAME'
 ATTRIBUTE_ERROR_MESSAGE = 'ERROR_MESSAGE'
 RECV_PREFIX = 'Recv'
 
-RpcRequestInfo = collections.namedtuple(
-    'RPCRequestInfo', ('request', 'context')
-)
-# exc_info is the three tuple defined at:
-# https://docs.python.org/3/library/sys.html#sys.exc_info
-RpcResponseInfo = collections.namedtuple(
-    'RPCCallbackInfo', ('request', 'context', 'response', 'exc_info')
-)
-
-
-class RpcMethodHandlerWrapper(object):
-    """Wraps a grpc RPCMethodHandler and records the variables about the
-     execution context and response
-    """
-
-    def __init__(
-        self, handler, pre_handler_callbacks=None, post_handler_callbacks=None
-    ):
-        """
-        :param handler: instance of RpcMethodHandler
-
-        :param pre_handler_callbacks: iterable of callbacks that accept an
-         instance of RpcRequestInfo that are called before the server handler
-
-        :param post_handler_callbacks: iterable of callbacks that accept an
-         instance of RpcResponseInfo that are called after the server
-         handler finishes execution
-        """
-        self.handler = handler
-        self._pre_handler_callbacks = pre_handler_callbacks or []
-        self._post_handler_callbacks = post_handler_callbacks or []
-
-    def proxy(self, prop_name):
-        def _wrapper(request, context, *args, **kwargs):
-            for callback in self._pre_handler_callbacks:
-                callback(RpcRequestInfo(request, context))
-            exc_info = (None, None, None)
-            response = None
-            try:
-                response = getattr(
-                    self.handler, prop_name
-                )(request, context, *args, **kwargs)
-            except Exception as e:
-                logging.exception(e)
-                exc_info = sys.exc_info()
-                raise
-            finally:
-                for callback in self._post_handler_callbacks:
-                    callback(RpcResponseInfo(
-                        request, context, response, exc_info)
-                    )
-            return response
-
-        return _wrapper
-
-    def __getattr__(self, item):
-        if item in (
-            'unary_unary', 'unary_stream', 'stream_unary', 'stream_stream'
-        ):
-            return self.proxy(item)
-        return getattr(self.handler, item)
-
 
 class OpenCensusServerInterceptor(grpc.ServerInterceptor):
     def __init__(self, sampler=None, exporter=None):
         self.sampler = sampler
         self.exporter = exporter
 
-    def _start_server_span(self, rpc_request_info):
-        metadata = rpc_request_info.context.invocation_metadata()
+    def intercept_service(self, continuation, handler_call_details):
+        def trace_wrapper(behavior, request_streaming, response_streaming):
+            def new_behavior(request_or_iterator, servicer_context):
+                span = self._start_server_span(servicer_context)
+                try:
+                    if request_streaming:
+                        request_or_iterator = grpc_utils.wrap_iter_with_message_events(  # noqa: E501
+                            request_or_response_iter=request_or_iterator,
+                            span=span,
+                            message_event_type=time_event.Type.RECEIVED
+                        )
+                    else:
+                        grpc_utils.add_message_event(
+                            proto_message=request_or_iterator,
+                            span=span,
+                            message_event_type=time_event.Type.RECEIVED,
+                        )
+                    # invoke the original rpc behavior
+                    response_or_iterator = behavior(request_or_iterator,
+                                                    servicer_context)
+                    if response_streaming:
+                        response_or_iterator = grpc_utils.wrap_iter_with_message_events(  # noqa: E501
+                            request_or_response_iter=response_or_iterator,
+                            span=span,
+                            message_event_type=time_event.Type.SENT
+                        )
+                        response_or_iterator = grpc_utils.wrap_iter_with_end_span(  # noqa: E501
+                            response_or_iterator)
+                    else:
+                        grpc_utils.add_message_event(
+                            proto_message=response_or_iterator,
+                            span=span,
+                            message_event_type=time_event.Type.SENT,
+                        )
+                except Exception as exc:
+                    logging.exception(exc)
+                    _add_exc_info(span)
+                    raise
+                finally:
+                    # if the response is unary, end the span here. Otherwise
+                    # it will be closed when the response iter completes
+                    if not response_streaming:
+                        execution_context.get_opencensus_tracer().end_span()
+                return response_or_iterator
+
+            return new_behavior
+
+        return _wrap_rpc_behavior(
+            continuation(handler_call_details),
+            trace_wrapper
+        )
+
+    def _start_server_span(self, servicer_context):
+        metadata = servicer_context.invocation_metadata()
         span_context = None
 
         if metadata is not None:
@@ -114,7 +105,9 @@ class OpenCensusServerInterceptor(grpc.ServerInterceptor):
                                       sampler=self.sampler,
                                       exporter=self.exporter)
 
-        span = tracer.start_span(name=_get_span_name(rpc_request_info))
+        span = tracer.start_span(
+            name=_get_span_name(servicer_context)
+        )
         tracer.add_attribute_to_current_span(
             attribute_key=attributes_helper.COMMON_ATTRIBUTES.get(
                 ATTRIBUTE_COMPONENT),
@@ -122,41 +115,52 @@ class OpenCensusServerInterceptor(grpc.ServerInterceptor):
 
         execution_context.set_opencensus_tracer(tracer)
         execution_context.set_current_span(span)
-
-    def _end_server_span(self, rpc_response_info):
-        tracer = execution_context.get_opencensus_tracer()
-        exc_type, exc_value, tb = rpc_response_info.exc_info
-        if exc_type is not None:
-            current_span = tracer.current_span()
-            current_span.add_attribute(
-                attributes_helper.COMMON_ATTRIBUTES.get(
-                    ATTRIBUTE_ERROR_MESSAGE),
-                str(exc_value)
-            )
-            current_span.stack_trace = stack_trace.StackTrace.from_traceback(
-                tb
-            )
-            current_span.status = status.Status(
-                code=code_pb2.UNKNOWN,
-                message=str(exc_value)
-            )
-
-        tracer.end_span()
-
-    def intercept_handler(self, continuation, handler_call_details):
-        return RpcMethodHandlerWrapper(
-            continuation(handler_call_details),
-            pre_handler_callbacks=[self._start_server_span],
-            post_handler_callbacks=[self._end_server_span]
-        )
-
-    def intercept_service(self, continuation, handler_call_details):
-        return self.intercept_handler(continuation, handler_call_details)
-
-
-def _get_span_name(rpc_request_info):
+        return span
+
+
+def _add_exc_info(span):
+    exc_type, exc_value, tb = sys.exc_info()
+    span.add_attribute(
+        attributes_helper.COMMON_ATTRIBUTES.get(
+            ATTRIBUTE_ERROR_MESSAGE),
+        str(exc_value)
+    )
+    span.stack_trace = stack_trace.StackTrace.from_traceback(tb)
+    span.status = status.Status(
+        code=code_pb2.UNKNOWN,
+        message=str(exc_value)
+    )
+
+
+def _wrap_rpc_behavior(handler, fn):
+    """Returns a new rpc handler that wraps the given function"""
+    if handler is None:
+        return None
+
+    if handler.request_streaming and handler.response_streaming:
+        behavior_fn = handler.stream_stream
+        handler_factory = grpc.stream_stream_rpc_method_handler
+    elif handler.request_streaming and not handler.response_streaming:
+        behavior_fn = handler.stream_unary
+        handler_factory = grpc.stream_unary_rpc_method_handler
+    elif not handler.request_streaming and handler.response_streaming:
+        behavior_fn = handler.unary_stream
+        handler_factory = grpc.unary_stream_rpc_method_handler
+    else:
+        behavior_fn = handler.unary_unary
+        handler_factory = grpc.unary_unary_rpc_method_handler
+
+    return handler_factory(
+        fn(behavior_fn, handler.request_streaming,
+           handler.response_streaming),
+        request_deserializer=handler.request_deserializer,
+        response_serializer=handler.response_serializer
+    )
+
+
+def _get_span_name(servicer_context):
     """Generates a span name based off of the gRPC server rpc_request_info"""
-    method_name = rpc_request_info.context._rpc_event.call_details.method[1:]
+    method_name = servicer_context._rpc_event.call_details.method[1:]
     if isinstance(method_name, bytes):
         method_name = method_name.decode('utf-8')
     method_name = method_name.replace('/', '.')
diff --git a/opencensus/trace/ext/grpc/utils.py b/opencensus/trace/ext/grpc/utils.py
new file mode 100644
index 00000000..7028d4f
--- /dev/null
+++ b/opencensus/trace/ext/grpc/utils.py
@@ -0,0 +1,45 @@
+from datetime import datetime
+
+from opencensus.trace import time_event
+from opencensus.trace import execution_context
+
+
+def add_message_event(proto_message, span, message_event_type, message_id=1):
+    """Adds a MessageEvent to the span based off of the given protobuf
+    message
+    """
+    span.add_time_event(
+        time_event=time_event.TimeEvent(
+            datetime.utcnow(),
+            message_event=time_event.MessageEvent(
+                message_id,
+                type=message_event_type,
+                uncompressed_size_bytes=proto_message.ByteSize()
+            )
+        )
+    )
+
+
+def wrap_iter_with_message_events(
+    request_or_response_iter,
+    span,
+    message_event_type
+):
+    """Wraps a request or response iterator to add message events to the span
+    for each proto message sent or received
+    """
+    for message_id, message in enumerate(request_or_response_iter, start=1):
+        add_message_event(
+            proto_message=message,
+            span=span,
+            message_event_type=message_event_type,
+            message_id=message_id
+        )
+        yield message
+
+
+def wrap_iter_with_end_span(response_iter):
+    """Wraps an iterator to end the current span on completion"""
+    for response in response_iter:
+        yield response
+    execution_context.get_opencensus_tracer().end_span()


Test patch:
diff --git a/tests/unit/trace/ext/grpc/test_client_interceptor.py b/tests/unit/trace/ext/grpc/test_client_interceptor.py
index 617c477..0783318 100644
--- a/tests/unit/trace/ext/grpc/test_client_interceptor.py
+++ b/tests/unit/trace/ext/grpc/test_client_interceptor.py
@@ -15,15 +15,13 @@
 import unittest
 
 import mock
-
+from opencensus.trace import execution_context
 from opencensus.trace.ext.grpc import client_interceptor
 
 
 class TestOpenCensusClientInterceptor(unittest.TestCase):
 
     def setUp(self):
-        from opencensus.trace import execution_context
-
         execution_context.clear()
 
     def test_constructor_default(self):
@@ -183,11 +181,11 @@ class TestOpenCensusClientInterceptor(unittest.TestCase):
 
     def _stream_helper(self):
         continuation = mock.Mock()
-        mock_response = mock.Mock()
+        mock_response = iter([mock.Mock()])
         continuation.return_value = mock_response
         mock_tracer = mock.Mock()
         interceptor = client_interceptor.OpenCensusClientInterceptor(tracer=mock_tracer)
-        interceptor._intercept_call = mock.Mock(return_value=(None, iter([mock.Mock()]), None))
+        interceptor._intercept_call = mock.Mock(return_value=(None, iter([mock.Mock()]), mock.Mock()))
         return interceptor, continuation, mock_tracer
 
     def test_intercept_unary_unary_trace(self):
@@ -207,9 +205,14 @@ class TestOpenCensusClientInterceptor(unittest.TestCase):
 
     def test_intercept_unary_stream_trace(self):
         interceptor, continuation, mock_tracer = self._stream_helper()
+        execution_context.set_opencensus_tracer(mock_tracer)
         client_call_details = mock.Mock()
         client_call_details.method = 'test'
-        interceptor.intercept_unary_stream(continuation, client_call_details, [])
+        response_iter = interceptor.intercept_unary_stream(
+            continuation, client_call_details, []
+        )
+        for _ in response_iter:
+            pass
         self.assertTrue(mock_tracer.end_span.called)
 
     def test_intercept_unary_stream_not_trace(self):
@@ -237,9 +240,15 @@ class TestOpenCensusClientInterceptor(unittest.TestCase):
 
     def test_intercept_stream_stream_trace(self):
         interceptor, continuation, mock_tracer = self._stream_helper()
+        execution_context.set_opencensus_tracer(mock_tracer)
         client_call_details = mock.Mock()
         client_call_details.method = 'test'
-        interceptor.intercept_stream_stream(continuation, client_call_details, [])
+        response_iter = interceptor.intercept_stream_stream(
+            continuation, client_call_details, []
+        )
+        for _ in response_iter:
+            pass
+
         self.assertTrue(mock_tracer.end_span.called)
 
     def test_intercept_stream_stream_not_trace(self):
diff --git a/tests/unit/trace/ext/grpc/test_server_interceptor.py b/tests/unit/trace/ext/grpc/test_server_interceptor.py
index 1565def..abd21ce 100644
--- a/tests/unit/trace/ext/grpc/test_server_interceptor.py
+++ b/tests/unit/trace/ext/grpc/test_server_interceptor.py
@@ -31,14 +31,7 @@ class TestOpenCensusServerInterceptor(unittest.TestCase):
         self.assertEqual(interceptor.sampler, sampler)
         self.assertEqual(interceptor.exporter, exporter)
 
-    def test_rpc_handler_wrapper(self):
-        """Ensure that RPCHandlerWrapper proxies to the unerlying handler"""
-        mock_handler = mock.Mock()
-        mock_handler.response_streaming = False
-        wrapper = server_interceptor.RpcMethodHandlerWrapper(mock_handler)
-        self.assertEqual(wrapper.response_streaming, False)
-
-    def test_intercept_handler_no_metadata(self):
+    def test_intercept_service_no_metadata(self):
         patch = mock.patch(
             'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
             MockTracer)
@@ -47,35 +40,14 @@ class TestOpenCensusServerInterceptor(unittest.TestCase):
         mock_context._rpc_event.call_details.method = 'hello'
         interceptor = server_interceptor.OpenCensusServerInterceptor(
             None, None)
+        mock_handler = mock.Mock()
+        mock_handler.request_streaming = False
+        mock_handler.response_streaming = False
+        mock_continuation = mock.Mock(return_value=mock_handler)
 
         with patch:
-            interceptor.intercept_handler(
-                mock.Mock(), mock.Mock()
-            ).unary_unary(mock.Mock(), mock_context)
-
-        expected_attributes = {
-            '/component': 'grpc',
-        }
-
-        self.assertEqual(
-            execution_context.get_opencensus_tracer().current_span().attributes,
-            expected_attributes)
-
-    def test_intercept_handler(self):
-        patch = mock.patch(
-            'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
-            MockTracer)
-        mock_context = mock.Mock()
-        mock_context.invocation_metadata = mock.Mock(
-            return_value=(('test_key', b'test_value'),)
-        )
-        mock_context._rpc_event.call_details.method = 'hello'
-        interceptor = server_interceptor.OpenCensusServerInterceptor(
-            None, None)
-
-        with patch:
-            interceptor.intercept_handler(
-                mock.Mock(), mock.Mock()
+            interceptor.intercept_service(
+                mock_continuation, mock.Mock()
             ).unary_unary(mock.Mock(), mock_context)
 
         expected_attributes = {
@@ -87,52 +59,97 @@ class TestOpenCensusServerInterceptor(unittest.TestCase):
             expected_attributes)
 
     def test_intercept_service(self):
-        interceptor = server_interceptor.OpenCensusServerInterceptor(
-            None, None)
-        mock_handler = mock.Mock()
-        interceptor.intercept_handler = mock_handler
-        interceptor.intercept_service(None, None)
-        self.assertTrue(mock_handler.called)
+        test_dimensions = [
+            ['unary_unary', False, False],
+            ['unary_stream', False, True],
+            ['stream_unary', True, False],
+            ['stream_stream', True, True],
+        ]
+        for rpc_fn_name, request_streaming, response_streaming in test_dimensions:
+            patch = mock.patch(
+                'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
+                MockTracer)
+            mock_context = mock.Mock()
+            mock_context.invocation_metadata = mock.Mock(
+                return_value=(('test_key', b'test_value'),)
+            )
+            mock_handler = mock.Mock()
+            mock_handler.request_streaming = request_streaming
+            mock_handler.response_streaming = response_streaming
+            mock_continuation = mock.Mock(return_value=mock_handler)
+
+            mock_context._rpc_event.call_details.method = 'hello'
+            interceptor = server_interceptor.OpenCensusServerInterceptor(
+                None, None)
+
+            with patch:
+                handler = interceptor.intercept_service(
+                    mock_continuation, mock.Mock()
+                )
+                getattr(handler, rpc_fn_name)(mock.Mock(), mock_context)
+
+            expected_attributes = {
+                '/component': 'grpc',
+            }
+
+            self.assertEqual(
+                execution_context.get_opencensus_tracer().current_span().attributes,
+                expected_attributes)
 
     def test_intercept_handler_exception(self):
-        patch = mock.patch(
-            'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
-            MockTracer)
-        interceptor = server_interceptor.OpenCensusServerInterceptor(
-            None, None)
-        mock_context = mock.Mock()
-        mock_context.invocation_metadata = mock.Mock(return_value=None)
-        mock_context._rpc_event.call_details.method = 'hello'
-        mock_continuation = mock.Mock()
-        mock_continuation.unary_unary = mock.Mock(side_effect=Exception('Test'))
-        with patch:
-            # patch the wrapper's handler to return an exception
-            rpc_wrapper = interceptor.intercept_handler(
-                mock.Mock(), mock.Mock())
-            rpc_wrapper.handler.unary_unary = mock.Mock(
-                side_effect=Exception('Test'))
-            with self.assertRaises(Exception):
-                rpc_wrapper.unary_unary(mock.Mock(), mock_context)
-
-        expected_attributes = {
-            '/component': 'grpc',
-            '/error/message': 'Test'
-        }
-
-        current_span = execution_context.get_opencensus_tracer().current_span()
-        self.assertEqual(
-            execution_context.get_opencensus_tracer().current_span().attributes,
-            expected_attributes)
-
-        # check that the stack trace is attached to the current span
-        self.assertIsNotNone(current_span.stack_trace)
-        self.assertIsNotNone(current_span.stack_trace.stack_trace_hash_id)
-        self.assertNotEqual(current_span.stack_trace.stack_frames, [])
-
-        # check that the status obj is attached to the current span
-        self.assertIsNotNone(current_span.status)
-        self.assertEqual(current_span.status.code, code_pb2.UNKNOWN)
-        self.assertEqual(current_span.status.message, 'Test')
+        test_dimensions = [
+            ['unary_unary', False, False],
+            ['unary_stream', False, True],
+            ['stream_unary', True, False],
+            ['stream_stream', True, True],
+        ]
+        for rpc_fn_name, request_streaming, response_streaming in test_dimensions:
+            patch = mock.patch(
+                'opencensus.trace.ext.grpc.server_interceptor.tracer_module.Tracer',
+                MockTracer)
+            interceptor = server_interceptor.OpenCensusServerInterceptor(
+                None, None)
+            mock_context = mock.Mock()
+            mock_context.invocation_metadata = mock.Mock(return_value=None)
+            mock_context._rpc_event.call_details.method = 'hello'
+            mock_handler = mock.Mock()
+            mock_handler.request_streaming = request_streaming
+            mock_handler.response_streaming = response_streaming
+            setattr(mock_handler, rpc_fn_name,
+                    mock.Mock(side_effect=Exception('Test')))
+            mock_continuation = mock.Mock(return_value=mock_handler)
+
+            # mock_continuation.unary_unary = mock.Mock())
+            with patch:
+                # patch the wrapper's handler to return an exception
+                handler = interceptor.intercept_service(
+                    mock_continuation, mock.Mock())
+                with self.assertRaises(Exception):
+                    getattr(handler, rpc_fn_name)(mock.Mock(), mock_context)
+
+            expected_attributes = {
+                '/component': 'grpc',
+                '/error/message': 'Test'
+            }
+
+            current_span = execution_context.get_opencensus_tracer().current_span()
+            self.assertEqual(
+                execution_context.get_opencensus_tracer().current_span().attributes,
+                expected_attributes)
+
+            # check that the stack trace is attached to the current span
+            self.assertIsNotNone(current_span.stack_trace)
+            self.assertIsNotNone(current_span.stack_trace.stack_trace_hash_id)
+            self.assertNotEqual(current_span.stack_trace.stack_frames, [])
+
+            # check that the status obj is attached to the current span
+            self.assertIsNotNone(current_span.status)
+            self.assertEqual(current_span.status.code, code_pb2.UNKNOWN)
+            self.assertEqual(current_span.status.message, 'Test')
+
+    def test__wrap_rpc_behavior_none(self):
+        new_handler = server_interceptor._wrap_rpc_behavior(None, lambda: None)
+        self.assertEqual(new_handler, None)
 
 
 class MockTracer(object):