xref: /aosp_15_r20/external/pigweed/pw_rpc/py/pw_rpc/testing.py (revision 61c4878ac05f98d0ceed94b57d316916de578985) 
1# Copyright 2021 The Pigweed Authors
2#
3# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4# use this file except in compliance with the License. You may obtain a copy of
5# the License at
6#
7#     https://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, WITHOUT
11# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12# License for the specific language governing permissions and limitations under
13# the License.
14"""Utilities for testing pw_rpc."""
15
16import argparse
17import shlex
18import subprocess
19import sys
20import tempfile
21import time
22from typing import Sequence
23
24TEMP_DIR_MARKER = '(pw_rpc:CREATE_TEMP_DIR)'
25
26
27def parse_test_server_args(
28    parser: argparse.ArgumentParser | None = None,
29) -> argparse.Namespace:
30    """Parses arguments for running a Python-based integration test."""
31    if parser is None:
32        parser = argparse.ArgumentParser(
33            description=sys.modules['__main__'].__doc__
34        )
35
36    parser.add_argument(
37        '--test-server-command',
38        nargs='+',
39        required=True,
40        help='Command that starts the test server.',
41    )
42    parser.add_argument(
43        '--port',
44        type=int,
45        required=True,
46        help=(
47            'The port to use to connect to the test server. This value is '
48            'passed to the test server as the last argument.'
49        ),
50    )
51    parser.add_argument(
52        'unittest_args',
53        nargs=argparse.REMAINDER,
54        help='Arguments after "--" are passed to unittest.',
55    )
56
57    args = parser.parse_args()
58
59    # Append the port number to the test server command.
60    args.test_server_command.append(str(args.port))
61
62    # Make the script name argv[0] and drop the "--".
63    args.unittest_args = sys.argv[:1] + args.unittest_args[1:]
64
65    return args
66
67
68def _parse_subprocess_integration_test_args() -> argparse.Namespace:
69    parser = argparse.ArgumentParser(
70        description='Executes a test between two subprocesses'
71    )
72    parser.add_argument(
73        '--client',
74        required=True,
75        help=(
76            'Client command to run. '
77            'Use quotes and whitespace to pass client-specifc arguments.'
78        ),
79    )
80    parser.add_argument(
81        '--server',
82        required=True,
83        help=(
84            'Server command to run. '
85            'Use quotes and whitespace to pass client-specifc arguments.'
86        ),
87    )
88    parser.add_argument(
89        'common_args',
90        metavar='-- ...',
91        nargs=argparse.REMAINDER,
92        help=(
93            'Arguments to pass to both the server and client; '
94            f'pass {TEMP_DIR_MARKER} to generate a temporary directory'
95        ),
96    )
97
98    args = parser.parse_args()
99
100    if not args.common_args or args.common_args[0] != '--':
101        parser.error('The common arguments must start with "--"')
102
103    args.common_args.pop(0)
104
105    return args
106
107
108def execute_integration_test(
109    server: str,
110    client: str,
111    common_args: Sequence[str],
112    setup_time_s: float = 0.2,
113) -> int:
114    """Runs an RPC server and client as part of an integration test."""
115    temp_dir: tempfile.TemporaryDirectory | None = None
116
117    if TEMP_DIR_MARKER in common_args:
118        temp_dir = tempfile.TemporaryDirectory(prefix='pw_rpc_test_')
119        common_args = [
120            temp_dir.name if a == TEMP_DIR_MARKER else a for a in common_args
121        ]
122
123    try:
124        server_cmdline = shlex.split(server)
125        client_cmdline = shlex.split(client)
126        if common_args:
127            server_cmdline += [*common_args]
128            client_cmdline += [*common_args]
129
130        server_process = subprocess.Popen(server_cmdline)
131        # TODO: b/234879791 - Replace this delay with some sort of IPC.
132        time.sleep(setup_time_s)
133
134        result = subprocess.run(client_cmdline).returncode
135
136        server_process.terminate()
137        server_process.communicate()
138    finally:
139        if temp_dir:
140            temp_dir.cleanup()
141
142    return result
143
144
145if __name__ == '__main__':
146    sys.exit(
147        execute_integration_test(
148            **vars(_parse_subprocess_integration_test_args())
149        )
150    )
151