• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright 2020 The gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Test of propagation of contextvars to AuthMetadataPlugin threads.."""
15
16import contextlib
17import logging
18import os
19import sys
20import unittest
21
22import grpc
23
24from tests.unit import test_common
25
26_UNARY_UNARY = "/test/UnaryUnary"
27_REQUEST = b"0000"
28
29
30def _unary_unary_handler(request, context):
31    return request
32
33
34def contextvars_supported():
35    try:
36        import contextvars
37        return True
38    except ImportError:
39        return False
40
41
42class _GenericHandler(grpc.GenericRpcHandler):
43
44    def service(self, handler_call_details):
45        if handler_call_details.method == _UNARY_UNARY:
46            return grpc.unary_unary_rpc_method_handler(_unary_unary_handler)
47        else:
48            raise NotImplementedError()
49
50
51@contextlib.contextmanager
52def _server():
53    try:
54        server = test_common.test_server()
55        target = 'localhost:0'
56        port = server.add_insecure_port(target)
57        server.add_generic_rpc_handlers((_GenericHandler(),))
58        server.start()
59        yield port
60    finally:
61        server.stop(None)
62
63
64if contextvars_supported():
65    import contextvars
66
67    _EXPECTED_VALUE = 24601
68    test_var = contextvars.ContextVar("test_var", default=None)
69
70    def set_up_expected_context():
71        test_var.set(_EXPECTED_VALUE)
72
73    class TestCallCredentials(grpc.AuthMetadataPlugin):
74
75        def __call__(self, context, callback):
76            if test_var.get() != _EXPECTED_VALUE:
77                raise AssertionError("{} != {}".format(test_var.get(),
78                                                       _EXPECTED_VALUE))
79            callback((), None)
80
81        def assert_called(self, test):
82            test.assertTrue(self._invoked)
83            test.assertEqual(_EXPECTED_VALUE, self._recorded_value)
84
85else:
86
87    def set_up_expected_context():
88        pass
89
90    class TestCallCredentials(grpc.AuthMetadataPlugin):
91
92        def __call__(self, context, callback):
93            callback((), None)
94
95
96# TODO(https://github.com/grpc/grpc/issues/22257)
97@unittest.skipIf(os.name == "nt", "LocalCredentials not supported on Windows.")
98class ContextVarsPropagationTest(unittest.TestCase):
99
100    def test_propagation_to_auth_plugin(self):
101        set_up_expected_context()
102        with _server() as port:
103            target = "localhost:{}".format(port)
104            local_credentials = grpc.local_channel_credentials()
105            test_call_credentials = TestCallCredentials()
106            call_credentials = grpc.metadata_call_credentials(
107                test_call_credentials, "test call credentials")
108            composite_credentials = grpc.composite_channel_credentials(
109                local_credentials, call_credentials)
110            with grpc.secure_channel(target, composite_credentials) as channel:
111                stub = channel.unary_unary(_UNARY_UNARY)
112                response = stub(_REQUEST, wait_for_ready=True)
113                self.assertEqual(_REQUEST, response)
114
115
116if __name__ == '__main__':
117    logging.basicConfig()
118    unittest.main(verbosity=2)
119