wehub-resource-sync
140 lines
4.9 KiB
Python
140 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock
|
|
import json
|
|
|
|
# Third Party
|
|
from fastapi.testclient import TestClient
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.internal_api_server.api_server import app
|
|
|
|
|
|
def _get_test_scripts_static():
|
|
"""Static function to get test scripts for parametrize."""
|
|
script_dir = Path(__file__).parent / "test_scripts"
|
|
scripts = []
|
|
for script_file in sorted(script_dir.glob("*.py")):
|
|
content = script_file.read_text()
|
|
# Find TEST_METADATA line
|
|
lines = content.split("\n")
|
|
metadata_lines = []
|
|
in_metadata = False
|
|
brace_count = 0
|
|
|
|
for line in lines:
|
|
if line.strip().startswith("# TEST_METADATA: {"):
|
|
in_metadata = True
|
|
# Extract JSON part after ': '
|
|
json_part = line.split("# TEST_METADATA: ", 1)[1]
|
|
metadata_lines.append(json_part)
|
|
brace_count = json_part.count("{") - json_part.count("}")
|
|
elif in_metadata and line.strip().startswith("#"):
|
|
# Continuation line in comment
|
|
json_part = line.strip()[1:].strip() # Remove '#' and trim
|
|
metadata_lines.append(json_part)
|
|
brace_count += json_part.count("{") - json_part.count("}")
|
|
|
|
# Check if we've closed all braces
|
|
if brace_count == 0 and "}" in json_part:
|
|
in_metadata = False
|
|
elif in_metadata:
|
|
# Not a comment line, metadata block ended
|
|
in_metadata = False
|
|
|
|
if metadata_lines:
|
|
json_str = " ".join(metadata_lines)
|
|
metadata = json.loads(json_str)
|
|
scripts.append((script_file.name, metadata))
|
|
return scripts
|
|
|
|
|
|
class TestRunScriptAPI:
|
|
"""Test suite for the /run_script API endpoint."""
|
|
|
|
@pytest.fixture
|
|
def mock_lmcache_adapter_with_config(self):
|
|
"""Create a mock LMCacheConnectorV1Impl adapter with config."""
|
|
adapter = MagicMock()
|
|
config = LMCacheEngineConfig.from_defaults()
|
|
adapter.config = config
|
|
return adapter
|
|
|
|
@pytest.fixture
|
|
def client_with_adapter(self, mock_lmcache_adapter_with_config):
|
|
"""Create a test client with mocked adapter."""
|
|
app.state.lmcache_adapter = mock_lmcache_adapter_with_config
|
|
app.state.lmcache_engine = mock_lmcache_adapter_with_config.lmcache_engine
|
|
client = TestClient(app)
|
|
yield client
|
|
client.close()
|
|
|
|
@pytest.fixture
|
|
def script_dir(self):
|
|
"""Get the test scripts directory."""
|
|
return Path(__file__).parent / "test_scripts"
|
|
|
|
@pytest.mark.parametrize(
|
|
"script_name,metadata",
|
|
[
|
|
pytest.param(name, meta, id=name.replace(".py", ""))
|
|
for name, meta in _get_test_scripts_static()
|
|
],
|
|
)
|
|
def test_run_script_from_file(
|
|
self,
|
|
client_with_adapter,
|
|
script_dir,
|
|
script_name,
|
|
metadata,
|
|
):
|
|
"""Test script execution by directly uploading script files."""
|
|
script_path = script_dir / script_name
|
|
|
|
# Configure allowed imports
|
|
allowed_imports = metadata.get("allowed_imports", [])
|
|
client_with_adapter.app.state.lmcache_adapter.config.script_allowed_imports = (
|
|
allowed_imports
|
|
)
|
|
|
|
# Upload the script file directly
|
|
with open(script_path, "rb") as f:
|
|
files = {"script": (script_name, f, "text/plain")}
|
|
response = client_with_adapter.post("/run_script", files=files)
|
|
|
|
# Verify status code
|
|
expected_status = metadata.get("expected_status", 200)
|
|
assert response.status_code == expected_status, (
|
|
f"Expected status {expected_status}, got {response.status_code}"
|
|
)
|
|
|
|
# Verify response content
|
|
if "expected_result" in metadata:
|
|
assert response.text == metadata["expected_result"], (
|
|
f"Expected '{metadata['expected_result']}', got '{response.text}'"
|
|
)
|
|
|
|
if "expected_contains" in metadata:
|
|
contains_list = metadata["expected_contains"]
|
|
if isinstance(contains_list, str):
|
|
contains_list = [contains_list]
|
|
for expected_text in contains_list:
|
|
assert expected_text in response.text, (
|
|
f"Expected '{expected_text}' in response"
|
|
)
|
|
|
|
if "expected_min_length" in metadata:
|
|
assert len(response.text) >= metadata["expected_min_length"], (
|
|
f"Response too short: {len(response.text)}"
|
|
)
|
|
|
|
def test_run_script_no_file(self, client_with_adapter):
|
|
"""Test run_script with no file provided."""
|
|
response = client_with_adapter.post("/run_script")
|
|
|
|
assert response.status_code == 400
|
|
assert "No script file provided" in response.text
|