main
py 95 lines 3.11 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from unittest.mock import Mock
5
6 import pytest
7
8 import tests.utils as test_utils
9 from tests.utils import RemoteLaunchRenderServer, RemoteOpenAIServer
10
11
12 @pytest.mark.parametrize(
13 ("args", "expected"),
14 [
15 (["--api-key", "key1", "key2"], ["--api-key", "***", "***"]),
16 (["--api_key", "key1", "key2"], ["--api_key", "***", "***"]),
17 (["--api-key=key1", "key2"], ["--api-key=***", "***"]),
18 (["--api_key=key1", "key2"], ["--api_key=***", "***"]),
19 (["--hf-token", "token"], ["--hf-token", "***"]),
20 (["--hf_token", "token"], ["--hf_token", "***"]),
21 (["--hf-token=token"], ["--hf-token=***"]),
22 (["--hf_token=token"], ["--hf_token=***"]),
23 ],
24 )
25 def test_redact_sensitive_cli_arg_variants(args: list[str], expected: list[str]):
26 assert test_utils._redact_sensitive_cli_args(args) == expected
27
28
29 @pytest.mark.parametrize(
30 ("server_cls", "command"),
31 [
32 (RemoteOpenAIServer, ["vllm", "serve"]),
33 (RemoteLaunchRenderServer, ["vllm", "launch", "render"]),
34 ],
35 )
36 def test_server_redacts_sensitive_values(
37 server_cls,
38 command: list[str],
39 monkeypatch: pytest.MonkeyPatch,
40 capsys: pytest.CaptureFixture[str],
41 ):
42 server = object.__new__(server_cls)
43 popen = Mock()
44 monkeypatch.setattr(test_utils.subprocess, "Popen", popen)
45 monkeypatch.setenv("INHERITED_SECRET", "inherited-secret")
46 serve_args = [
47 "--api_key",
48 "api-secret-1",
49 "api-secret-2",
50 "--hf-token=hf-secret",
51 "--max-num-seqs",
52 "2",
53 ]
54
55 server._start_server(
56 "test-model", serve_args, {"OVERRIDE_SECRET": "override-secret"}
57 )
58
59 stdout = capsys.readouterr().out
60 assert "api-secret-1" not in stdout
61 assert "api-secret-2" not in stdout
62 assert "hf-secret" not in stdout
63 assert "inherited-secret" not in stdout
64 assert "override-secret" not in stdout
65 assert "--api_key *** ***" in stdout
66 assert "--hf-token=***" in stdout
67 assert "--max-num-seqs 2" in stdout
68 assert popen.call_args.args[0] == [*command, "test-model", *serve_args]
69 child_env = popen.call_args.kwargs["env"]
70 assert child_env["INHERITED_SECRET"] == "inherited-secret"
71 assert child_env["OVERRIDE_SECRET"] == "override-secret"
72
73
74 def test_openai_server_shutdown_wait_covers_engine_cleanup(
75 monkeypatch: pytest.MonkeyPatch,
76 ):
77 server = object.__new__(RemoteOpenAIServer)
78 server._request_shutdown_timeout = 0.0
79 server.proc = Mock(pid=1234)
80 engine_timeout_args = []
81
82 def get_engine_timeout(request_timeout, process_timeout):
83 engine_timeout_args.append((request_timeout, process_timeout))
84 return 60.0
85
86 monkeypatch.setattr(
87 test_utils, "get_engine_process_shutdown_timeout", get_engine_timeout
88 )
89 monkeypatch.setattr(test_utils.os, "getpgid", lambda pid: pid)
90 monkeypatch.setattr(server, "_kill_process_group_survivors", Mock())
91
92 server._terminate_process_tree()
93
94 assert engine_timeout_args == [(0.0, 0.0)]
95 server.proc.wait.assert_called_once_with(timeout=75.0)