AdamWen95
70 lines
3.0 KiB
Python
70 lines
3.0 KiB
Python
import json
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
from pathlib import Path
|
|
import tempfile
|
|
import threading
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from console.model_preflight import check_model
|
|
|
|
|
|
class ModelPreflightTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp.name)
|
|
self.credential = self.root / 'vllm.env'
|
|
self.credential.write_text('VLLM_API_KEY="dummy-metadata-key"\n')
|
|
self.calls = []
|
|
self.status = 200
|
|
self.body = {'data': [{'id': 'current-model', 'max_model_len': 262144}]}
|
|
case = self
|
|
class Handler(BaseHTTPRequestHandler):
|
|
def log_message(self, *args): pass
|
|
def do_GET(self):
|
|
case.calls.append((self.path, self.headers.get('Authorization')))
|
|
self.send_response(case.status)
|
|
if case.status == 302:
|
|
self.send_header('Location', '/redirect-target')
|
|
self.end_headers()
|
|
self.wfile.write(json.dumps(case.body).encode())
|
|
self.server = HTTPServer(('127.0.0.1', 0), Handler)
|
|
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
|
self.thread.start()
|
|
self.target = {'url': 'http://127.0.0.1:' + str(self.server.server_port), 'model': 'current-model'}
|
|
|
|
def tearDown(self):
|
|
self.server.shutdown(); self.server.server_close(); self.thread.join()
|
|
self.temp.cleanup()
|
|
|
|
def test_metadata_only_one_request_with_runtime_credentials(self):
|
|
result = check_model(self.target, [str(self.credential)])
|
|
self.assertTrue(result['ok'])
|
|
self.assertEqual(self.calls, [('/v1/models', 'Bearer dummy-metadata-key')])
|
|
self.assertNotIn('dummy-metadata-key', json.dumps(result))
|
|
|
|
def test_wrong_model_fails_even_when_port_is_open(self):
|
|
self.target['model'] = 'stale-model'
|
|
result = check_model(self.target, [str(self.credential)])
|
|
self.assertFalse(result['ok'])
|
|
self.assertIn('模型名称', result['message'])
|
|
|
|
def test_auth_failure_and_redirect_do_not_retry_or_leak(self):
|
|
for status in (401, 302):
|
|
self.status = status; self.calls.clear()
|
|
result = check_model(self.target, [str(self.credential)])
|
|
self.assertFalse(result['ok'])
|
|
self.assertEqual(len(self.calls), 1)
|
|
self.assertNotIn('dummy-metadata-key', json.dumps(result))
|
|
|
|
def test_external_address_rejected_before_network(self):
|
|
self.target['url'] = 'http://example.com'
|
|
with patch('console.model_preflight.urllib.request.build_opener') as build:
|
|
self.assertFalse(check_model(self.target, [])['ok'])
|
|
build.assert_not_called()
|
|
|
|
def test_malformed_or_oversize_metadata_is_not_success(self):
|
|
for body in ({'data': {}}, {'data': ['bad']}, {'data': [], 'padding': 'x' * 262145}):
|
|
self.body = body
|
|
self.assertFalse(check_model(self.target, [str(self.credential)])['ok'])
|