Use single import point for protobuf in grpc (#647)

Co-authored-by: Hannah Stepanek <[email protected]>

Author: lrafeei

tests/framework_grpc/_test_common.py

@@ -12,23 +12,23 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import grpc
-import threading
 import functools
+import threading
+
 from newrelic.api.application import application_instance
 
 
 def create_request(streaming_request, count=1, timesout=False):
-    from sample_application.sample_application_pb2 import Message
+    from sample_application import Message
 
     def _message_stream():
         for i in range(count):
-            yield Message(text='Hello World', count=count, timesout=timesout)
+            yield Message(text="Hello World", count=count, timesout=timesout)
 
     if streaming_request:
         request = _message_stream()
     else:
-        request = Message(text='Hello World', count=count, timesout=timesout)
+        request = Message(text="Hello World", count=count, timesout=timesout)
 
     return request
 

tests/framework_grpc/conftest.py

@@ -13,52 +13,56 @@
 # limitations under the License.
 
 import gc
+
 import grpc
 import pytest
-import random
-
-from testing_support.fixtures import (code_coverage_fixture,
-        collector_agent_registration_fixture, collector_available_fixture)
+from testing_support.fixtures import (  # noqa
+    code_coverage_fixture,
+    collector_agent_registration_fixture,
+    collector_available_fixture,
+)
 from testing_support.mock_external_grpc_server import MockExternalgRPCServer
+
 import newrelic.packages.six as six
 
 _coverage_source = [
-    'newrelic.hooks.framework_grpc',
+    "newrelic.hooks.framework_grpc",
 ]
 
 code_coverage = code_coverage_fixture(source=_coverage_source)
 
 _default_settings = {
-    'transaction_tracer.explain_threshold': 0.0,
-    'transaction_tracer.transaction_threshold': 0.0,
-    'transaction_tracer.stack_trace_threshold': 0.0,
-    'debug.log_data_collector_payloads': True,
-    'debug.record_transaction_failure': True,
+    "transaction_tracer.explain_threshold": 0.0,
+    "transaction_tracer.transaction_threshold": 0.0,
+    "transaction_tracer.stack_trace_threshold": 0.0,
+    "debug.log_data_collector_payloads": True,
+    "debug.record_transaction_failure": True,
 }
 
 collector_agent_registration = collector_agent_registration_fixture(
-        app_name='Python Agent Test (framework_grpc)',
-        default_settings=_default_settings)
+    app_name="Python Agent Test (framework_grpc)", default_settings=_default_settings
+)
 
 
-@pytest.fixture(scope='session')
+@pytest.fixture(scope="session")
 def grpc_app_server():
     with MockExternalgRPCServer() as server:
         yield server, server.port
 
 
-@pytest.fixture(scope='session')
+@pytest.fixture(scope="session")
 def mock_grpc_server(grpc_app_server):
-    from sample_application.sample_application_pb2_grpc import (
-            add_SampleApplicationServicer_to_server)
-    from sample_application import SampleApplicationServicer
+    from sample_application import (
+        SampleApplicationServicer,
+        add_SampleApplicationServicer_to_server,
+    )
+
     server, port = grpc_app_server
-    add_SampleApplicationServicer_to_server(
-            SampleApplicationServicer(), server)
+    add_SampleApplicationServicer_to_server(SampleApplicationServicer(), server)
     return port
 
 
-@pytest.fixture(scope='function', autouse=True)
+@pytest.fixture(scope="function", autouse=True)
 def gc_garbage_empty():
     yield
 
@@ -72,8 +76,8 @@ def gc_garbage_empty():
             pass
 
         from grpc._channel import _Rendezvous
-        rendezvous_stored = sum(1 for o in gc.get_objects()
-                if hasattr(o, '__class__') and isinstance(o, _Rendezvous))
+
+        rendezvous_stored = sum(1 for o in gc.get_objects() if hasattr(o, "__class__") and isinstance(o, _Rendezvous))
 
         assert rendezvous_stored == 0
 
@@ -89,17 +93,14 @@ def stub(stub_and_channel):
 @pytest.fixture(scope="session")
 def stub_and_channel(mock_grpc_server):
     port = mock_grpc_server
-    from sample_application.sample_application_pb2_grpc import (
-            SampleApplicationStub)
-
     stub, channel = create_stub_and_channel(port)
     with channel:
         yield stub, channel
 
+
 def create_stub_and_channel(port):
-    from sample_application.sample_application_pb2_grpc import (
-            SampleApplicationStub)
+    from sample_application import SampleApplicationStub
 
-    channel = grpc.insecure_channel('localhost:%s' % port)
+    channel = grpc.insecure_channel("localhost:%s" % port)
     stub = SampleApplicationStub(channel)
     return stub, channel

tests/framework_grpc/sample_application/__init__.py

@@ -14,84 +14,94 @@
 
 import json
 import time
-from newrelic.api.transaction import current_transaction
+
 import grpc
+import sample_application_pb2_grpc
 
-from sample_application_pb2 import Message
-from sample_application_pb2_grpc import (
-        SampleApplicationServicer as _SampleApplicationServicer)
+from newrelic.api.transaction import current_transaction
+
+# This import format is to resolve a bug within protobuf 4
+# Issues for reference:
+# https://github.com/protocolbuffers/protobuf/issues/10075
+# https://github.com/protocolbuffers/protobuf/issues/10151
+# Within sample_application_pb2.py, the protobuf import can only
+# be done once before the DESCRIPTOR value is set to None
+# (in subsequent imports) instead of overriding/ignoring the imports.
+# This ensures that the imports happen once.
+Message = sample_application_pb2_grpc.sample__application__pb2.Message
+add_SampleApplicationServicer_to_server = sample_application_pb2_grpc.add_SampleApplicationServicer_to_server
+SampleApplicationStub = sample_application_pb2_grpc.SampleApplicationStub
 
 
 class Status(object):
     code = grpc.StatusCode.ABORTED
-    details = 'abort_with_status'
+    details = "abort_with_status"
     trailing_metadata = {}
 
 
-class SampleApplicationServicer(_SampleApplicationServicer):
-
+class SampleApplicationServicer(sample_application_pb2_grpc.SampleApplicationServicer):
     def DoUnaryUnary(self, request, context):
-        context.set_trailing_metadata([('content-type', 'text/plain')])
+        context.set_trailing_metadata([("content-type", "text/plain")])
         if request.timesout:
             while context.is_active():
                 time.sleep(0.1)
-        return Message(text='unary_unary: %s' % request.text)
+        return Message(text="unary_unary: %s" % request.text)
 
     def DoUnaryStream(self, request, context):
-        context.set_trailing_metadata([('content-type', 'text/plain')])
+        context.set_trailing_metadata([("content-type", "text/plain")])
         if request.timesout:
             while context.is_active():
                 time.sleep(0.1)
         for i in range(request.count):
-            yield Message(text='unary_stream: %s' % request.text)
+            yield Message(text="unary_stream: %s" % request.text)
 
     def DoStreamUnary(self, request_iter, context):
-        context.set_trailing_metadata([('content-type', 'text/plain')])
+        context.set_trailing_metadata([("content-type", "text/plain")])
         for request in request_iter:
             if request.timesout:
                 while context.is_active():
                     time.sleep(0.1)
-            return Message(text='stream_unary: %s' % request.text)
+            return Message(text="stream_unary: %s" % request.text)
 
     def DoStreamStream(self, request_iter, context):
-        context.set_trailing_metadata([('content-type', 'text/plain')])
+        context.set_trailing_metadata([("content-type", "text/plain")])
         for request in request_iter:
             if request.timesout:
                 while context.is_active():
                     time.sleep(0.1)
-            yield Message(text='stream_stream: %s' % request.text)
+            yield Message(text="stream_stream: %s" % request.text)
 
     def DoUnaryUnaryRaises(self, request, context):
-        raise AssertionError('unary_unary: %s' % request.text)
+        raise AssertionError("unary_unary: %s" % request.text)
 
     def DoUnaryStreamRaises(self, request, context):
-        raise AssertionError('unary_stream: %s' % request.text)
+        raise AssertionError("unary_stream: %s" % request.text)
 
     def DoStreamUnaryRaises(self, request_iter, context):
         for request in request_iter:
-            raise AssertionError('stream_unary: %s' % request.text)
+            raise AssertionError("stream_unary: %s" % request.text)
 
     def DoStreamStreamRaises(self, request_iter, context):
         for request in request_iter:
-            raise AssertionError('stream_stream: %s' % request.text)
+            raise AssertionError("stream_stream: %s" % request.text)
 
     def NoTxnUnaryUnaryRaises(self, request, context):
         current_transaction().ignore_transaction = True
-        raise AssertionError('unary_unary: %s' % request.text)
+        raise AssertionError("unary_unary: %s" % request.text)
 
     def NoTxnUnaryStreamRaises(self, request, context):
         current_transaction().ignore_transaction = True
-        raise AssertionError('unary_stream: %s' % request.text)
+        raise AssertionError("unary_stream: %s" % request.text)
 
     def NoTxnStreamUnaryRaises(self, request_iter, context):
         current_transaction().ignore_transaction = True
         for request in request_iter:
-            raise AssertionError('stream_unary: %s' % request.text)
+            raise AssertionError("stream_unary: %s" % request.text)
 
     def NoTxnStreamStreamRaises(self, request_iter, context):
         current_transaction().ignore_transaction = True
         for request in request_iter:
-            raise AssertionError('stream_stream: %s' % request.text)
+            raise AssertionError("stream_stream: %s" % request.text)
 
     def NoTxnUnaryUnary(self, request, context):
         current_transaction().ignore_transaction = True
@@ -110,16 +120,16 @@ def NoTxnStreamStream(self, request_iter, context):
         return self.DoStreamStream(request_iter, context)
 
     def DoUnaryUnaryAbort(self, request, context):
-        context.abort(grpc.StatusCode.ABORTED, 'aborting')
+        context.abort(grpc.StatusCode.ABORTED, "aborting")
 
     def DoUnaryStreamAbort(self, request, context):
-        context.abort(grpc.StatusCode.ABORTED, 'aborting')
+        context.abort(grpc.StatusCode.ABORTED, "aborting")
 
     def DoStreamUnaryAbort(self, request_iter, context):
-        context.abort(grpc.StatusCode.ABORTED, 'aborting')
+        context.abort(grpc.StatusCode.ABORTED, "aborting")
 
     def DoStreamStreamAbort(self, request_iter, context):
-        context.abort(grpc.StatusCode.ABORTED, 'aborting')
+        context.abort(grpc.StatusCode.ABORTED, "aborting")
 
     def DoUnaryUnaryAbortWithStatus(self, request, context):
         context.abort_with_status(Status)

tox.ini

@@ -133,7 +133,7 @@ envlist =
     python-framework_graphql-{py37,py38,py39,py310,pypy37}-graphql03,
     ; temporarily disabling graphqlmaster tests
     python-framework_graphql-py37-graphql{0202,0203,0300,0301,0302},
-    grpc-framework_grpc-{py27}-grpc0125,
+    grpc-framework_grpc-py27-grpc0125,
     grpc-framework_grpc-{py37,py38,py39,py310}-grpclatest,
     python-framework_pyramid-{pypy,py27,py38}-Pyramid0104,
     python-framework_pyramid-{pypy,py27,pypy37,py37,py38,py39,py310}-Pyramid0110-cornice,
@@ -320,7 +320,7 @@ deps =
     framework_graphql-graphql0301: graphql-core<3.2
     framework_graphql-graphql0302: graphql-core<3.3
     framework_graphql-graphqlmaster: https://github.com/graphql-python/graphql-core/archive/main.zip
-    framework_grpc-grpclatest: protobuf<4
+    framework_grpc-grpclatest: protobuf
     framework_grpc-grpclatest: grpcio
     framework_grpc-grpclatest: grpcio-tools
     grpc0125: grpcio<1.26