_signal_handling_test.py

Go to the documentation of this file.

# Copyright 2019 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.
"""Test of responsiveness to signals."""
 
import logging
import os
import signal
import subprocess
import sys
import tempfile
import threading
import unittest
 
import grpc
 
from tests.unit import _signal_client
from tests.unit import test_common
 
_CLIENT_PATH = None
if sys.executable is not None:
    _CLIENT_PATH = os.path.abspath(os.path.realpath(_signal_client.__file__))
else:
    # NOTE(rbellevi): For compatibility with internal testing.
    if len(sys.argv) != 2:
        raise RuntimeError("Must supply path to executable client.")
    client_name = sys.argv[1].split("/")[-1]
    del sys.argv[1]  # For compatibility with test runner.
    _CLIENT_PATH = os.path.realpath(
        os.path.join(os.path.dirname(os.path.abspath(__file__)), client_name))
 
_HOST = 'localhost'
 
# The gevent test harness cannot run the monkeypatch code for the child process,
# so we need to instrument it manually.
_GEVENT_ARG = ("--gevent",) if test_common.running_under_gevent() else ()
 
 
class _GenericHandler(grpc.GenericRpcHandler):
 
    def __init__(self):
        self._connected_clients_lock = threading.RLock()
        self._connected_clients_event = threading.Event()
        self._connected_clients = 0
 
        self._unary_unary_handler = grpc.unary_unary_rpc_method_handler(
            self._handle_unary_unary)
        self._unary_stream_handler = grpc.unary_stream_rpc_method_handler(
            self._handle_unary_stream)
 
    def _on_client_connect(self):
        with self._connected_clients_lock:
            self._connected_clients += 1
            self._connected_clients_event.set()
 
    def _on_client_disconnect(self):
        with self._connected_clients_lock:
            self._connected_clients -= 1
            if self._connected_clients == 0:
                self._connected_clients_event.clear()
 
    def await_connected_client(self):
        """Blocks until a client connects to the server."""
        self._connected_clients_event.wait()
 
    def _handle_unary_unary(self, request, servicer_context):
        """Handles a unary RPC.
 
        Blocks until the client disconnects and then echoes.
        """
        stop_event = threading.Event()
 
        def on_rpc_end():
            self._on_client_disconnect()
            stop_event.set()
 
        servicer_context.add_callback(on_rpc_end)
        self._on_client_connect()
        stop_event.wait()
        return request
 
    def _handle_unary_stream(self, request, servicer_context):
        """Handles a server streaming RPC.
 
        Blocks until the client disconnects and then echoes.
        """
        stop_event = threading.Event()
 
        def on_rpc_end():
            self._on_client_disconnect()
            stop_event.set()
 
        servicer_context.add_callback(on_rpc_end)
        self._on_client_connect()
        stop_event.wait()
        yield request
 
    def service(self, handler_call_details):
        if handler_call_details.method == _signal_client.UNARY_UNARY:
            return self._unary_unary_handler
        elif handler_call_details.method == _signal_client.UNARY_STREAM:
            return self._unary_stream_handler
        else:
            return None
 
 
def _read_stream(stream):
    stream.seek(0)
    return stream.read()
 
 
def _start_client(args, stdout, stderr):
    invocation = None
    if sys.executable is not None:
        invocation = (sys.executable, _CLIENT_PATH) + tuple(args)
    else:
        invocation = (_CLIENT_PATH,) + tuple(args)
    return subprocess.Popen(invocation, stdout=stdout, stderr=stderr)
 
 
class SignalHandlingTest(unittest.TestCase):
 
    def setUp(self):
        self._server = test_common.test_server()
        self._port = self._server.add_insecure_port('{}:0'.format(_HOST))
        self._handler = _GenericHandler()
        self._server.add_generic_rpc_handlers((self._handler,))
        self._server.start()
 
    def tearDown(self):
        self._server.stop(None)
 
    @unittest.skipIf(os.name == 'nt', 'SIGINT not supported on windows')
    def testUnary(self):
        """Tests that the server unary code path does not stall signal handlers."""
        server_target = '{}:{}'.format(_HOST, self._port)
        with tempfile.TemporaryFile(mode='r') as client_stdout:
            with tempfile.TemporaryFile(mode='r') as client_stderr:
                client = _start_client((server_target, 'unary') + _GEVENT_ARG,
                                       client_stdout, client_stderr)
                self._handler.await_connected_client()
                client.send_signal(signal.SIGINT)
                self.assertFalse(client.wait(), msg=_read_stream(client_stderr))
                client_stdout.seek(0)
                self.assertIn(_signal_client.SIGTERM_MESSAGE,
                              client_stdout.read())
 
    @unittest.skipIf(os.name == 'nt', 'SIGINT not supported on windows')
    def testStreaming(self):
        """Tests that the server streaming code path does not stall signal handlers."""
        server_target = '{}:{}'.format(_HOST, self._port)
        with tempfile.TemporaryFile(mode='r') as client_stdout:
            with tempfile.TemporaryFile(mode='r') as client_stderr:
                client = _start_client(
                    (server_target, 'streaming') + _GEVENT_ARG, client_stdout,
                    client_stderr)
                self._handler.await_connected_client()
                client.send_signal(signal.SIGINT)
                self.assertFalse(client.wait(), msg=_read_stream(client_stderr))
                client_stdout.seek(0)
                self.assertIn(_signal_client.SIGTERM_MESSAGE,
                              client_stdout.read())
 
    @unittest.skipIf(os.name == 'nt', 'SIGINT not supported on windows')
    def testUnaryWithException(self):
        server_target = '{}:{}'.format(_HOST, self._port)
        with tempfile.TemporaryFile(mode='r') as client_stdout:
            with tempfile.TemporaryFile(mode='r') as client_stderr:
                client = _start_client(
                    ('--exception', server_target, 'unary') + _GEVENT_ARG,
                    client_stdout, client_stderr)
                self._handler.await_connected_client()
                client.send_signal(signal.SIGINT)
                client.wait()
                self.assertEqual(0, client.returncode)
 
    @unittest.skipIf(os.name == 'nt', 'SIGINT not supported on windows')
    def testStreamingHandlerWithException(self):
        server_target = '{}:{}'.format(_HOST, self._port)
        with tempfile.TemporaryFile(mode='r') as client_stdout:
            with tempfile.TemporaryFile(mode='r') as client_stderr:
                client = _start_client(
                    ('--exception', server_target, 'streaming') + _GEVENT_ARG,
                    client_stdout, client_stderr)
                self._handler.await_connected_client()
                client.send_signal(signal.SIGINT)
                client.wait()
                print(_read_stream(client_stderr))
                self.assertEqual(0, client.returncode)
 
 
if __name__ == '__main__':
    logging.basicConfig()
    unittest.main(verbosity=2)