From f13b3249414240ce0a28110ba1a7f6732a6c3ec8 Mon Sep 17 00:00:00 2001
From: Richard Belleville <rbellevi@google.com>
Date: Mon, 19 Oct 2020 13:02:33 -0700
Subject: [PATCH] Include well-known types in sys.path when using runtime
 protos

---
 .../grpcio_tests/tests/unit/BUILD.bazel       |  1 +
 .../tests/unit/_dynamic_stubs_test.py         | 47 ++++++++++++++-----
 .../tests/unit/data/foo/baz.proto             | 27 +++++++++++
 tools/distrib/python/grpcio_tools/BUILD.bazel |  1 +
 .../python/grpcio_tools/grpc_tools/protoc.py  |  2 +
 5 files changed, 66 insertions(+), 12 deletions(-)
 create mode 100644 src/python/grpcio_tests/tests/unit/data/foo/baz.proto

diff --git a/src/python/grpcio_tests/tests/unit/BUILD.bazel b/src/python/grpcio_tests/tests/unit/BUILD.bazel
index 909ced21383..64dcab29615 100644
--- a/src/python/grpcio_tests/tests/unit/BUILD.bazel
+++ b/src/python/grpcio_tests/tests/unit/BUILD.bazel
@@ -125,6 +125,7 @@ py2and3_test(
     srcs = ["_dynamic_stubs_test.py"],
     data = [
         "data/foo/bar.proto",
+        "data/foo/baz.proto",
     ],
     imports = ["../../"],
     main = "_dynamic_stubs_test.py",
diff --git a/src/python/grpcio_tests/tests/unit/_dynamic_stubs_test.py b/src/python/grpcio_tests/tests/unit/_dynamic_stubs_test.py
index 6a4387c547c..ed107742b67 100644
--- a/src/python/grpcio_tests/tests/unit/_dynamic_stubs_test.py
+++ b/src/python/grpcio_tests/tests/unit/_dynamic_stubs_test.py
@@ -21,6 +21,8 @@ import os
 import sys
 import unittest
 
+_DATA_DIR = os.path.join("tests", "unit", "data")
+
 
 @contextlib.contextmanager
 def _grpc_tools_unimportable():
@@ -53,6 +55,18 @@ def _collect_errors(fn):
     return _wrapped
 
 
+def _python3_check(fn):
+
+    @functools.wraps(fn)
+    def _wrapped():
+        if sys.version_info[0] == 3:
+            fn()
+        else:
+            _assert_unimplemented("Python 3")
+
+    return _wrapped
+
+
 def _run_in_subprocess(test_case):
     sys.path.insert(
         0, os.path.join(os.path.realpath(os.path.dirname(__file__)), ".."))
@@ -80,24 +94,30 @@ def _assert_unimplemented(msg_substr):
 
 
 @_collect_errors
+@_python3_check
 def _test_sunny_day():
-    if sys.version_info[0] == 3:
-        import grpc
-        protos, services = grpc.protos_and_services(
-            os.path.join("tests", "unit", "data", "foo", "bar.proto"))
-        assert protos.BarMessage is not None
-        assert services.BarStub is not None
-    else:
-        _assert_unimplemented("Python 3")
+    import grpc
+    protos, services = grpc.protos_and_services(
+        os.path.join(_DATA_DIR, "foo", "bar.proto"))
+    assert protos.BarMessage is not None
+    assert services.BarStub is not None
 
 
 @_collect_errors
+@_python3_check
+def _test_well_known_types():
+    import grpc
+    protos, services = grpc.protos_and_services(
+        os.path.join(_DATA_DIR, "foo", "baz.proto"))
+    assert protos.BarMessage is not None
+    assert services.BarStub is not None
+
+
+@_collect_errors
+@_python3_check
 def _test_grpc_tools_unimportable():
     with _grpc_tools_unimportable():
-        if sys.version_info[0] == 3:
-            _assert_unimplemented("grpcio-tools")
-        else:
-            _assert_unimplemented("Python 3")
+        _assert_unimplemented("grpcio-tools")
 
 
 # NOTE(rbellevi): multiprocessing.Process fails to pickle function objects
@@ -109,6 +129,9 @@ class DynamicStubTest(unittest.TestCase):
     def test_sunny_day(self):
         _run_in_subprocess(_test_sunny_day)
 
+    def test_well_known_types(self):
+        _run_in_subprocess(_test_well_known_types)
+
     def test_grpc_tools_unimportable(self):
         _run_in_subprocess(_test_grpc_tools_unimportable)
 
diff --git a/src/python/grpcio_tests/tests/unit/data/foo/baz.proto b/src/python/grpcio_tests/tests/unit/data/foo/baz.proto
new file mode 100644
index 00000000000..a2ca7ab524f
--- /dev/null
+++ b/src/python/grpcio_tests/tests/unit/data/foo/baz.proto
@@ -0,0 +1,27 @@
+// Copyright 2020 The gRPC 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.
+
+syntax = "proto3";
+
+package tests.unit.data.foo.bar;
+
+import "google/protobuf/wrappers.proto";
+
+message BarMessage {
+  string a = 1;
+};
+
+service Bar {
+  rpc GetBar(BarMessage) returns (BarMessage);
+};
diff --git a/tools/distrib/python/grpcio_tools/BUILD.bazel b/tools/distrib/python/grpcio_tools/BUILD.bazel
index c0b3f66e903..a23975feeb8 100644
--- a/tools/distrib/python/grpcio_tools/BUILD.bazel
+++ b/tools/distrib/python/grpcio_tools/BUILD.bazel
@@ -50,4 +50,5 @@ py_library(
         "//src/python/grpcio/grpc:grpcio",
         "@com_google_protobuf//:protobuf_python",
     ],
+    data = glob(["grpc_tools/_proto/**/*"]),
 )
diff --git a/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py b/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py
index d07daba2c06..48f853ad4f2 100644
--- a/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py
+++ b/tools/distrib/python/grpcio_tools/grpc_tools/protoc.py
@@ -58,6 +58,8 @@ if sys.version_info >= (3, 5, 0):
                     ProtoFinder(_SERVICE_MODULE_SUFFIX,
                                 _protoc_compiler.get_services)
                 ])
+                sys.path.append(
+                    pkg_resources.resource_filename('grpc_tools', '_proto'))
                 _FINDERS_INSTALLED = True
 
     def _module_name_to_proto_file(suffix, module_name):