wehub-resource-sync
103 lines
2.5 KiB
Python
103 lines
2.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import asyncio
|
|
import tempfile
|
|
from pathlib import Path
|
|
from ssl import SSLContext
|
|
|
|
import pytest
|
|
|
|
from vllm.entrypoints.serve.utils.ssl import SSLCertRefresher
|
|
|
|
|
|
class MockSSLContext(SSLContext):
|
|
def __init__(self):
|
|
self.load_cert_chain_count = 0
|
|
self.load_ca_count = 0
|
|
|
|
def load_cert_chain(
|
|
self,
|
|
certfile,
|
|
keyfile=None,
|
|
password=None,
|
|
):
|
|
self.load_cert_chain_count += 1
|
|
|
|
def load_verify_locations(
|
|
self,
|
|
cafile=None,
|
|
capath=None,
|
|
cadata=None,
|
|
):
|
|
self.load_ca_count += 1
|
|
|
|
|
|
def create_file() -> str:
|
|
with tempfile.NamedTemporaryFile(dir="/tmp", delete=False) as f:
|
|
return f.name
|
|
|
|
|
|
def touch_file(path: str) -> None:
|
|
Path(path).touch()
|
|
|
|
|
|
async def wait_for_counts(
|
|
ssl_context: MockSSLContext,
|
|
*,
|
|
cert_chain_count: int,
|
|
ca_count: int,
|
|
timeout: float = 5.0,
|
|
) -> None:
|
|
deadline = asyncio.get_running_loop().time() + timeout
|
|
while True:
|
|
if (
|
|
ssl_context.load_cert_chain_count >= cert_chain_count
|
|
and ssl_context.load_ca_count >= ca_count
|
|
):
|
|
return
|
|
|
|
if asyncio.get_running_loop().time() >= deadline:
|
|
assert ssl_context.load_cert_chain_count >= cert_chain_count
|
|
assert ssl_context.load_ca_count >= ca_count
|
|
|
|
await asyncio.sleep(0.05)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ssl_refresher():
|
|
ssl_context = MockSSLContext()
|
|
key_path = create_file()
|
|
cert_path = create_file()
|
|
ca_path = create_file()
|
|
ssl_refresher = SSLCertRefresher(ssl_context, key_path, cert_path, ca_path)
|
|
await asyncio.sleep(1)
|
|
assert ssl_context.load_cert_chain_count == 0
|
|
assert ssl_context.load_ca_count == 0
|
|
|
|
touch_file(key_path)
|
|
await wait_for_counts(
|
|
ssl_context,
|
|
cert_chain_count=1,
|
|
ca_count=0,
|
|
)
|
|
assert ssl_context.load_ca_count == 0
|
|
|
|
touch_file(cert_path)
|
|
touch_file(ca_path)
|
|
await wait_for_counts(
|
|
ssl_context,
|
|
cert_chain_count=2,
|
|
ca_count=1,
|
|
)
|
|
|
|
ssl_refresher.stop()
|
|
await asyncio.sleep(0)
|
|
cert_chain_count = ssl_context.load_cert_chain_count
|
|
ca_count = ssl_context.load_ca_count
|
|
|
|
touch_file(cert_path)
|
|
touch_file(ca_path)
|
|
await asyncio.sleep(1)
|
|
assert ssl_context.load_cert_chain_count == cert_chain_count
|
|
assert ssl_context.load_ca_count == ca_count
|