add test_grpc

mishards/grpc_utils/test_grpc.py

@@ -0,0 +1,77 @@
+import logging
+import opentracing
+from mishards.grpc_utils import GrpcSpanDecorator, is_grpc_method
+from milvus.grpc_gen import status_pb2, milvus_pb2
+
+
+logger = logging.getLogger(__name__)
+
+
+class TestTracer(opentracing.Tracer):
+    pass
+
+class TestSpan(opentracing.Span):
+    def __init__(self, context, tracer, **kwargs):
+        super(TestSpan, self).__init__(tracer, context)
+        self.reset()
+
+    def set_tag(self, key, value):
+        self.tags.append({key:value})
+
+    def log_kv(self, key_values, timestamp=None):
+        self.logs.append(key_values)
+
+    def reset(self):
+        self.tags = []
+        self.logs = []
+
+
+class TestRpcInfo:
+    def __init__(self, request, response):
+        self.request = request
+        self.response = response
+
+
+class TestGrpcUtils:
+    def test_span_deco(self):
+        request = 'request'
+        OK = status_pb2.Status(error_code=status_pb2.SUCCESS, reason='Success')
+        response = OK
+        rpc_info = TestRpcInfo(request=request, response=response)
+        span = TestSpan(context=None, tracer=TestTracer())
+        span_deco = GrpcSpanDecorator()
+        span_deco(span, rpc_info)
+        assert len(span.logs) == 0
+        assert len(span.tags) == 0
+
+        response = milvus_pb2.BoolReply(status=OK, bool_reply=False)
+        rpc_info = TestRpcInfo(request=request, response=response)
+        span = TestSpan(context=None, tracer=TestTracer())
+        span_deco = GrpcSpanDecorator()
+        span_deco(span, rpc_info)
+        assert len(span.logs) == 0
+        assert len(span.tags) == 0
+
+        response = 1
+        rpc_info = TestRpcInfo(request=request, response=response)
+        span = TestSpan(context=None, tracer=TestTracer())
+        span_deco = GrpcSpanDecorator()
+        span_deco(span, rpc_info)
+        logger.error(span.logs)
+        assert len(span.logs) == 1
+        assert len(span.tags) == 1
+
+        response = 0
+        rpc_info = TestRpcInfo(request=request, response=response)
+        span = TestSpan(context=None, tracer=TestTracer())
+        span_deco = GrpcSpanDecorator()
+        span_deco(span, rpc_info)
+        logger.error(span.logs)
+        assert len(span.logs) == 0
+        assert len(span.tags) == 0
+
+    def test_is_grpc_method(self):
+        target = 1
+        assert not is_grpc_method(target)
+        target = None
+        assert not is_grpc_method(target)