1 files changed, 73 insertions, 0 deletions
diff --git a/src/python/grpcio_test/grpc_test/beta/_beta_features_test.py b/src/python/grpcio_test/grpc_test/beta/_beta_features_test.py
index 5916a9e3ea..6b5090f8bb 100644
--- a/src/python/grpcio_test/grpc_test/beta/_beta_features_test.py
+++ b/src/python/grpcio_test/grpc_test/beta/_beta_features_test.py
@@ -224,5 +224,78 @@ class BetaFeaturesTest(unittest.TestCase):
     self.assertEqual(_RESPONSE, response)
 
 
+class ContextManagementAndLifecycleTest(unittest.TestCase):
+
+  def setUp(self):
+    self._servicer = _Servicer()
+    self._method_implementations = {
+        (_GROUP, _UNARY_UNARY):
+            utilities.unary_unary_inline(self._servicer.unary_unary),
+        (_GROUP, _UNARY_STREAM):
+            utilities.unary_stream_inline(self._servicer.unary_stream),
+        (_GROUP, _STREAM_UNARY):
+            utilities.stream_unary_inline(self._servicer.stream_unary),
+        (_GROUP, _STREAM_STREAM):
+            utilities.stream_stream_inline(self._servicer.stream_stream),
+    }
+
+    self._cardinalities = {
+        _UNARY_UNARY: cardinality.Cardinality.UNARY_UNARY,
+        _UNARY_STREAM: cardinality.Cardinality.UNARY_STREAM,
+        _STREAM_UNARY: cardinality.Cardinality.STREAM_UNARY,
+        _STREAM_STREAM: cardinality.Cardinality.STREAM_STREAM,
+    }
+
+    self._server_options = implementations.server_options(
+        thread_pool_size=test_constants.POOL_SIZE)
+    self._server_credentials = implementations.ssl_server_credentials(
+        [(resources.private_key(), resources.certificate_chain(),),])
+    self._client_credentials = implementations.ssl_client_credentials(
+        resources.test_root_certificates(), None, None)
+    self._stub_options = implementations.stub_options(
+        thread_pool_size=test_constants.POOL_SIZE)
+
+  def test_stub_context(self):
+    server = implementations.server(
+        self._method_implementations, options=self._server_options)
+    port = server.add_secure_port('[::]:0', self._server_credentials)
+    server.start()
+
+    channel = test_utilities.not_really_secure_channel(
+        'localhost', port, self._client_credentials, _SERVER_HOST_OVERRIDE)
+    dynamic_stub = implementations.dynamic_stub(
+        channel, _GROUP, self._cardinalities, options=self._stub_options)
+    for _ in range(100):
+      with dynamic_stub:
+        pass
+    for _ in range(10):
+      with dynamic_stub:
+        call_options = interfaces.grpc_call_options(
+            disable_compression=True)
+        response = getattr(dynamic_stub, _UNARY_UNARY)(
+            _REQUEST, test_constants.LONG_TIMEOUT,
+            protocol_options=call_options)
+        self.assertEqual(_RESPONSE, response)
+        self.assertIsNotNone(self._servicer.peer())
+
+    server.stop(test_constants.SHORT_TIMEOUT).wait()
+
+  def test_server_lifecycle(self):
+    for _ in range(100):
+      server = implementations.server(
+          self._method_implementations, options=self._server_options)
+      port = server.add_secure_port('[::]:0', self._server_credentials)
+      server.start()
+      server.stop(test_constants.SHORT_TIMEOUT).wait()
+    for _ in range(100):
+      server = implementations.server(
+          self._method_implementations, options=self._server_options)
+      server.add_secure_port('[::]:0', self._server_credentials)
+      server.add_insecure_port('[::]:0')
+      with server:
+        server.stop(test_constants.SHORT_TIMEOUT)
+      server.stop(test_constants.SHORT_TIMEOUT)
+
+
 if __name__ == '__main__':
   unittest.main(verbosity=2)