xref: /aosp_15_r20/external/grpc-grpc/test/cpp/naming/utils/health_check.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2) 
1#!/usr/bin/env python3
2# Copyright 2023 The gRPC Authors.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16import argparse
17import platform
18import subprocess
19import sys
20import time
21
22
23def test_runner_log(msg):
24    sys.stderr.write("\n%s: %s\n" % (__file__, msg))
25
26
27def python_args(arg_list):
28    if platform.system() == "Windows" and arg_list[0].endswith(".py"):
29        return [sys.executable] + arg_list
30    return arg_list
31
32
33def wait_until_dns_server_is_up(args):
34    for i in range(0, 30):
35        test_runner_log(
36            "Health check: attempt to connect to DNS server over TCP."
37        )
38        tcp_connect_subprocess = subprocess.Popen(
39            python_args(
40                [
41                    args.tcp_connect_bin_path,
42                    "--server_host",
43                    "127.0.0.1",
44                    "--server_port",
45                    str(args.dns_server_port),
46                    "--timeout",
47                    str(1),
48                ]
49            )
50        )
51        tcp_connect_subprocess.communicate()
52        if tcp_connect_subprocess.returncode == 0:
53            test_runner_log(
54                (
55                    "Health check: attempt to make an A-record "
56                    "query to DNS server."
57                )
58            )
59            dns_resolver_subprocess = subprocess.Popen(
60                python_args(
61                    [
62                        args.dns_resolver_bin_path,
63                        "--qname",
64                        "health-check-local-dns-server-is-alive.resolver-tests.grpctestingexp",
65                        "--server_host",
66                        "127.0.0.1",
67                        "--server_port",
68                        str(args.dns_server_port),
69                    ]
70                ),
71                stdout=subprocess.PIPE,
72            )
73            dns_resolver_stdout, _ = dns_resolver_subprocess.communicate(
74                str.encode("ascii")
75            )
76            if dns_resolver_subprocess.returncode == 0:
77                if "123.123.123.123".encode("ascii") in dns_resolver_stdout:
78                    test_runner_log(
79                        (
80                            "DNS server is up! "
81                            "Successfully reached it over UDP and TCP."
82                        )
83                    )
84                return
85        time.sleep(1)
86    test_runner_log(
87        (
88            "Failed to reach DNS server over TCP and/or UDP. "
89            "Exitting without running tests."
90        )
91    )
92    sys.exit(1)
93
94
95def main():
96    argp = argparse.ArgumentParser(description="Make DNS queries for A records")
97    argp.add_argument(
98        "-p",
99        "--dns_server_port",
100        default=None,
101        type=int,
102        help=("Port that local DNS server is listening on."),
103    )
104    argp.add_argument(
105        "--dns_resolver_bin_path",
106        default=None,
107        type=str,
108        help=("Path to the DNS health check utility."),
109    )
110    argp.add_argument(
111        "--tcp_connect_bin_path",
112        default=None,
113        type=str,
114        help=("Path to the TCP health check utility."),
115    )
116    args = argp.parse_args()
117    wait_until_dns_server_is_up(args)
118
119
120if __name__ == "__main__":
121    main()
122