Fix tests

Signed-off-by: Aidan Jensen <[email protected]>

Author: Aidan Jensen

instrumentation/opentelemetry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/_client.py

@@ -217,18 +217,17 @@ def _intercept_server_stream(
             stream = invoker(request_or_iterator, metadata)
 
             def done_callback(future, span_):
-                logger.exception("done_callback")
                 try:
                     future.result()
                 except grpc.FutureCancelledError:
                     span_.set_status(Status(StatusCode.OK))
                     span_.set_attribute(
-                        SpanAttributes.RPC_GRPC_STATUS_CODE, grpc.StatusCode.CANCELLED.value[0]
+                        RPC_GRPC_STATUS_CODE, grpc.StatusCode.CANCELLED.value[0]
                     )
                 except grpc.RpcError as err:
                     span_.set_status(Status(StatusCode.ERROR))
                     span_.set_attribute(
-                        SpanAttributes.RPC_GRPC_STATUS_CODE, err.code().value[0]
+                        RPC_GRPC_STATUS_CODE, err.code().value[0]
                     )
                 finally:
                     span_.end()

instrumentation/opentelemetry-instrumentation-grpc/tests/_client.py

@@ -57,16 +57,18 @@ def server_streaming_method(stub, error=False, serialize=True):
     return response_iterator
 
 
-def bidirectional_streaming_method(stub, error=False):
+def bidirectional_streaming_method(stub, error=False, serialize=True):
     def request_messages():
         for _ in range(5):
             request = Request(
                 client_id=CLIENT_ID, request_data="error" if error else "data"
             )
             yield request
 
-    response_iterator = stub.ServerStreamingMethod(
-        request, metadata=(("key", "value"),)
+    response_iterator = stub.BidirectionalStreamingMethod(
+        request_messages(), metadata=(("key", "value"),)
     )
 
-    list(response_iterator)
+    if serialize:
+        list(response_iterator)
+    return response_iterator

instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py

@@ -13,8 +13,9 @@
 # limitations under the License.
 # pylint:disable=cyclic-import
 import logging
-from time import sleep
 
+import threading
+import time
 from unittest import mock
 
 import grpc
@@ -174,16 +175,15 @@ def test_unary_stream(self):
         )
 
     def test_unary_stream_can_be_cancel(self):
+        done = threading.Event()
         responses = server_streaming_method(self._stub, serialize=False)
+        responses.add_done_callback(lambda: done.set())
         for i, _ in enumerate(responses):
             if i == 1:
                 responses.cancel()
                 break
         self.assertEqual(responses.code(), grpc.StatusCode.CANCELLED)
-        # self.server.stop(None)
-        # self.channel.close()
-        logging.exception("Getting spans")
-        # sleep(10)
+        done.wait(5)
         spans = self.memory_exporter.get_finished_spans()
         self.assertEqual(len(spans), 1)
         span = spans[0]
@@ -192,17 +192,17 @@ def test_unary_stream_can_be_cancel(self):
         self.assertIs(span.kind, trace.SpanKind.CLIENT)
 
         # Check version and name in span's instrumentation info
-        self.assertEqualSpanInstrumentationInfo(
+        self.assertEqualSpanInstrumentationScope(
             span, opentelemetry.instrumentation.grpc
         )
 
         self.assertSpanHasAttributes(
             span,
             {
-                SpanAttributes.RPC_METHOD: "ServerStreamingMethod",
-                SpanAttributes.RPC_SERVICE: "GRPCTestServer",
-                SpanAttributes.RPC_SYSTEM: "grpc",
-                SpanAttributes.RPC_GRPC_STATUS_CODE: grpc.StatusCode.CANCELLED.value[
+                RPC_METHOD: "ServerStreamingMethod",
+                RPC_SERVICE: "GRPCTestServer",
+                RPC_SYSTEM: "grpc",
+                RPC_GRPC_STATUS_CODE: grpc.StatusCode.CANCELLED.value[
                     0
                 ],
             },
@@ -258,6 +258,41 @@ def test_stream_stream(self):
             },
         )
 
+    def test_stream_stream_can_be_cancel(self):
+        done = threading.Event()
+        responses = bidirectional_streaming_method(self._stub, serialize=False)
+        responses.add_done_callback(lambda: done.set())
+        for i, _ in enumerate(responses):
+            if i == 1:
+                responses.cancel()
+                break
+        self.assertEqual(responses.code(), grpc.StatusCode.CANCELLED)
+        done.wait(5)
+        spans = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(spans), 1)
+        span = spans[0]
+
+        self.assertEqual(span.name, "/GRPCTestServer/BidirectionalStreamingMethod")
+        self.assertIs(span.kind, trace.SpanKind.CLIENT)
+
+        # Check version and name in span's instrumentation info
+        self.assertEqualSpanInstrumentationScope(
+            span, opentelemetry.instrumentation.grpc
+        )
+
+        self.assertSpanHasAttributes(
+            span,
+            {
+                RPC_METHOD: "BidirectionalStreamingMethod",
+                RPC_SERVICE: "GRPCTestServer",
+                RPC_SYSTEM: "grpc",
+                RPC_GRPC_STATUS_CODE: grpc.StatusCode.CANCELLED.value[
+                    0
+                ],
+            },
+        )
+
+
     def test_error_simple(self):
         with self.assertRaises(grpc.RpcError):
             simple_method(self._stub, error=True)