xref: /aosp_15_r20/external/grpc-grpc/examples/python/cancellation/test/_cancellation_example_test.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1# Copyright 2019 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 for cancellation example."""
15
16import contextlib
17import os
18import signal
19import socket
20import subprocess
21import unittest
22
23_BINARY_DIR = os.path.realpath(
24    os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
25)
26_SERVER_PATH = os.path.join(_BINARY_DIR, "server")
27_CLIENT_PATH = os.path.join(_BINARY_DIR, "client")
28
29
30@contextlib.contextmanager
31def _get_port():
32    sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
33    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
34    if sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT) == 0:
35        raise RuntimeError("Failed to set SO_REUSEPORT.")
36    sock.bind(("", 0))
37    try:
38        yield sock.getsockname()[1]
39    finally:
40        sock.close()
41
42
43def _start_client(
44    server_port, desired_string, ideal_distance, interesting_distance=None
45):
46    interesting_distance_args = (
47        ()
48        if interesting_distance is None
49        else ("--show-inferior", interesting_distance)
50    )
51    return subprocess.Popen(
52        (
53            _CLIENT_PATH,
54            desired_string,
55            "--server",
56            "localhost:{}".format(server_port),
57            "--ideal-distance",
58            str(ideal_distance),
59        )
60        + interesting_distance_args
61    )
62
63
64class CancellationExampleTest(unittest.TestCase):
65    def test_successful_run(self):
66        with _get_port() as test_port:
67            server_process = subprocess.Popen(
68                (_SERVER_PATH, "--port", str(test_port))
69            )
70            try:
71                client_process = _start_client(test_port, "aa", 0)
72                client_return_code = client_process.wait()
73                self.assertEqual(0, client_return_code)
74                self.assertIsNone(server_process.poll())
75            finally:
76                server_process.kill()
77                server_process.wait()
78
79    def test_graceful_sigint(self):
80        with _get_port() as test_port:
81            server_process = subprocess.Popen(
82                (_SERVER_PATH, "--port", str(test_port))
83            )
84            try:
85                client_process1 = _start_client(test_port, "aaaaaaaaaa", 0)
86                client_process1.send_signal(signal.SIGINT)
87                client_process1.wait()
88                client_process2 = _start_client(test_port, "aa", 0)
89                client_return_code = client_process2.wait()
90                self.assertEqual(0, client_return_code)
91                self.assertIsNone(server_process.poll())
92            finally:
93                server_process.kill()
94                server_process.wait()
95
96
97if __name__ == "__main__":
98    unittest.main(verbosity=2)
99