wehub-resource-sync
335 lines
13 KiB
Python
335 lines
13 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from unittest.mock import MagicMock
|
|
import json
|
|
|
|
# Third Party
|
|
from fastapi.testclient import TestClient
|
|
import pytest
|
|
|
|
# First Party
|
|
from lmcache.v1.internal_api_server.api_server import app
|
|
|
|
|
|
def _make_metadata(role: str) -> MagicMock:
|
|
"""Create a mock metadata with the given role."""
|
|
metadata = MagicMock()
|
|
metadata.role = role
|
|
return metadata
|
|
|
|
|
|
class TestLookupAPI:
|
|
"""Test suite for the /lookup/* API endpoints."""
|
|
|
|
@pytest.fixture
|
|
def mock_scheduler_manager(self):
|
|
"""Create a mock LMCacheManager for scheduler role."""
|
|
manager = MagicMock()
|
|
manager.lmcache_engine_metadata = _make_metadata("scheduler")
|
|
manager.get_lookup_info.return_value = {
|
|
"client": "HitLimitLookupClient(LMCacheBypassLookupClient)",
|
|
"server": "None",
|
|
}
|
|
manager.close_lookup_client.return_value = {
|
|
"old": "HitLimitLookupClient(LMCacheBypassLookupClient)"
|
|
}
|
|
manager.create_lookup_client.return_value = {"new": "LMCacheLookupClient"}
|
|
manager.recreate_lookup_client.return_value = {
|
|
"old": "HitLimitLookupClient(LMCacheBypassLookupClient)",
|
|
"new": "LMCacheLookupClient",
|
|
}
|
|
return manager
|
|
|
|
@pytest.fixture
|
|
def mock_worker_manager(self):
|
|
"""Create a mock LMCacheManager for worker role."""
|
|
manager = MagicMock()
|
|
manager.lmcache_engine_metadata = _make_metadata("worker")
|
|
manager.get_lookup_info.return_value = {
|
|
"client": "None",
|
|
"server": "LMCacheLookupServer",
|
|
}
|
|
manager.close_lookup_server.return_value = {"old": "LMCacheLookupServer"}
|
|
manager.create_lookup_server.return_value = {"new": "LMCacheAsyncLookupServer"}
|
|
manager.recreate_lookup_server.return_value = {
|
|
"old": "LMCacheLookupServer",
|
|
"new": "LMCacheAsyncLookupServer",
|
|
}
|
|
return manager
|
|
|
|
@pytest.fixture
|
|
def scheduler_client(self, mock_scheduler_manager):
|
|
"""Create a test client with mocked scheduler manager."""
|
|
app.state.lmcache_adapter = mock_scheduler_manager
|
|
client = TestClient(app)
|
|
yield client
|
|
client.close()
|
|
|
|
@pytest.fixture
|
|
def worker_client(self, mock_worker_manager):
|
|
"""Create a test client with mocked worker manager."""
|
|
app.state.lmcache_adapter = mock_worker_manager
|
|
client = TestClient(app)
|
|
yield client
|
|
client.close()
|
|
|
|
# ==================== GET /lookup/info Tests ====================
|
|
|
|
def test_get_lookup_info_scheduler(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test get lookup info for scheduler role."""
|
|
response = scheduler_client.get("/lookup/info")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert "LMCacheBypassLookupClient" in data["client"]
|
|
assert data["server"] == "None"
|
|
mock_scheduler_manager.get_lookup_info.assert_called_once()
|
|
|
|
def test_get_lookup_info_worker(self, worker_client, mock_worker_manager):
|
|
"""Test get lookup info for worker role."""
|
|
response = worker_client.get("/lookup/info")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["client"] == "None"
|
|
assert data["server"] == "LMCacheLookupServer"
|
|
mock_worker_manager.get_lookup_info.assert_called_once()
|
|
|
|
def test_get_lookup_info_not_supported(self):
|
|
"""Test get lookup info when manager doesn't support it."""
|
|
|
|
class SimpleAdapter:
|
|
pass
|
|
|
|
app.state.lmcache_adapter = SimpleAdapter()
|
|
with TestClient(app) as client:
|
|
response = client.get("/lookup/info")
|
|
|
|
assert response.status_code == 503
|
|
data = json.loads(response.text)
|
|
assert "error" in data
|
|
|
|
# ==================== POST /lookup/close Tests ====================
|
|
|
|
def test_close_lookup_scheduler(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test close lookup client for scheduler."""
|
|
response = scheduler_client.post("/lookup/close")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert "LMCacheBypassLookupClient" in data["old"]
|
|
assert data["role"] == "scheduler"
|
|
mock_scheduler_manager.close_lookup_client.assert_called_once()
|
|
|
|
def test_close_lookup_worker(self, worker_client, mock_worker_manager):
|
|
"""Test close lookup server for worker."""
|
|
response = worker_client.post("/lookup/close")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["old"] == "LMCacheLookupServer"
|
|
assert data["role"] == "worker"
|
|
mock_worker_manager.close_lookup_server.assert_called_once()
|
|
|
|
# ==================== POST /lookup/create Tests ====================
|
|
|
|
def test_create_lookup_scheduler(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test create lookup client for scheduler."""
|
|
response = scheduler_client.post("/lookup/create")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["new"] == "LMCacheLookupClient"
|
|
assert data["role"] == "scheduler"
|
|
mock_scheduler_manager.create_lookup_client.assert_called_once_with(
|
|
dryrun=False
|
|
)
|
|
|
|
def test_create_lookup_worker(self, worker_client, mock_worker_manager):
|
|
"""Test create lookup server for worker."""
|
|
response = worker_client.post("/lookup/create")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["new"] == "LMCacheAsyncLookupServer"
|
|
assert data["role"] == "worker"
|
|
mock_worker_manager.create_lookup_server.assert_called_once_with(dryrun=False)
|
|
|
|
def test_create_lookup_dryrun_scheduler(
|
|
self, scheduler_client, mock_scheduler_manager
|
|
):
|
|
"""Test dryrun create lookup client for scheduler."""
|
|
mock_scheduler_manager.create_lookup_client.return_value = {
|
|
"new": "LMCacheLookupClient",
|
|
"dryrun": True,
|
|
}
|
|
|
|
response = scheduler_client.post("/lookup/create?dryrun=true")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["new"] == "LMCacheLookupClient"
|
|
assert data["dryrun"] is True
|
|
mock_scheduler_manager.create_lookup_client.assert_called_once_with(dryrun=True)
|
|
|
|
def test_create_lookup_dryrun_worker(self, worker_client, mock_worker_manager):
|
|
"""Test dryrun create lookup server for worker."""
|
|
mock_worker_manager.create_lookup_server.return_value = {
|
|
"new": "LMCacheAsyncLookupServer",
|
|
"dryrun": True,
|
|
}
|
|
|
|
response = worker_client.post("/lookup/create?dryrun=true")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["new"] == "LMCacheAsyncLookupServer"
|
|
assert data["dryrun"] is True
|
|
mock_worker_manager.create_lookup_server.assert_called_once_with(dryrun=True)
|
|
|
|
def test_create_lookup_error(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test create lookup with error."""
|
|
mock_scheduler_manager.create_lookup_client.return_value = {
|
|
"error": "metadata not available"
|
|
}
|
|
|
|
response = scheduler_client.post("/lookup/create")
|
|
|
|
assert response.status_code == 400
|
|
data = json.loads(response.text)
|
|
assert data["error"] == "metadata not available"
|
|
|
|
# ==================== POST /lookup/recreate Tests ====================
|
|
|
|
def test_recreate_lookup_scheduler(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test recreate lookup client for scheduler."""
|
|
response = scheduler_client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert "LMCacheBypassLookupClient" in data["old"]
|
|
assert data["new"] == "LMCacheLookupClient"
|
|
assert data["role"] == "scheduler"
|
|
mock_scheduler_manager.recreate_lookup_client.assert_called_once()
|
|
|
|
def test_recreate_lookup_worker(self, worker_client, mock_worker_manager):
|
|
"""Test recreate lookup server for worker."""
|
|
response = worker_client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["old"] == "LMCacheLookupServer"
|
|
assert data["new"] == "LMCacheAsyncLookupServer"
|
|
assert data["role"] == "worker"
|
|
mock_worker_manager.recreate_lookup_server.assert_called_once()
|
|
|
|
def test_recreate_lookup_error(self, scheduler_client, mock_scheduler_manager):
|
|
"""Test recreate lookup with error."""
|
|
mock_scheduler_manager.recreate_lookup_client.return_value = {
|
|
"error": "only supported for scheduler role"
|
|
}
|
|
|
|
response = scheduler_client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 400
|
|
data = json.loads(response.text)
|
|
assert "error" in data
|
|
|
|
def test_recreate_lookup_unknown_role(self):
|
|
"""Test recreation with unknown role."""
|
|
manager = MagicMock()
|
|
manager.lmcache_engine_metadata = _make_metadata("unknown")
|
|
app.state.lmcache_adapter = manager
|
|
with TestClient(app) as client:
|
|
response = client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 400
|
|
data = json.loads(response.text)
|
|
assert data["error"] == "Unknown role"
|
|
|
|
def test_recreate_lookup_no_metadata(self):
|
|
"""Test recreation when metadata is not available."""
|
|
manager = MagicMock()
|
|
manager.lmcache_engine_metadata = None
|
|
app.state.lmcache_adapter = manager
|
|
with TestClient(app) as client:
|
|
response = client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 400
|
|
data = json.loads(response.text)
|
|
assert data["error"] == "Unknown role"
|
|
|
|
def test_recreate_lookup_not_supported(self):
|
|
"""Test recreation when manager doesn't support it."""
|
|
|
|
class SimpleAdapter:
|
|
lmcache_engine_metadata = _make_metadata("scheduler")
|
|
|
|
app.state.lmcache_adapter = SimpleAdapter()
|
|
with TestClient(app) as client:
|
|
response = client.post("/lookup/recreate")
|
|
|
|
assert response.status_code == 503
|
|
data = json.loads(response.text)
|
|
assert "error" in data
|
|
|
|
# ==================== Integration Flow Tests ====================
|
|
|
|
def test_full_flow_worker_then_scheduler(
|
|
self, mock_worker_manager, mock_scheduler_manager
|
|
):
|
|
"""
|
|
Test the complete flow:
|
|
1. POST /lookup/recreate on worker (recreate server)
|
|
2. POST /lookup/recreate on scheduler (recreate client)
|
|
|
|
Note: Config should be updated via /conf API before each step.
|
|
"""
|
|
# Step 1: Worker side - recreate server
|
|
app.state.lmcache_adapter = mock_worker_manager
|
|
with TestClient(app) as worker_client:
|
|
response = worker_client.post("/lookup/recreate")
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["old"] == "LMCacheLookupServer"
|
|
assert data["new"] == "LMCacheAsyncLookupServer"
|
|
mock_worker_manager.recreate_lookup_server.assert_called_once()
|
|
|
|
# Step 2: Scheduler side - recreate client
|
|
app.state.lmcache_adapter = mock_scheduler_manager
|
|
with TestClient(app) as scheduler_client:
|
|
response = scheduler_client.post("/lookup/recreate")
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert "LMCacheBypassLookupClient" in data["old"]
|
|
assert data["new"] == "LMCacheLookupClient"
|
|
mock_scheduler_manager.recreate_lookup_client.assert_called_once()
|
|
|
|
def test_step_by_step_flow(self, mock_scheduler_manager):
|
|
"""
|
|
Test step-by-step flow: dryrun -> close -> create.
|
|
"""
|
|
app.state.lmcache_adapter = mock_scheduler_manager
|
|
with TestClient(app) as client:
|
|
# Step 1: Dryrun - check what would be created
|
|
mock_scheduler_manager.create_lookup_client.return_value = {
|
|
"new": "LMCacheLookupClient",
|
|
"dryrun": True,
|
|
}
|
|
response = client.post("/lookup/create?dryrun=true")
|
|
assert response.status_code == 200
|
|
data = json.loads(response.text)
|
|
assert data["dryrun"] is True
|
|
|
|
# Step 2: Close current client
|
|
response = client.post("/lookup/close")
|
|
assert response.status_code == 200
|
|
mock_scheduler_manager.close_lookup_client.assert_called_once()
|
|
|
|
# Step 3: Create new client
|
|
mock_scheduler_manager.create_lookup_client.return_value = {
|
|
"new": "LMCacheLookupClient"
|
|
}
|
|
response = client.post("/lookup/create")
|
|
assert response.status_code == 200
|