mirror of
https://github.com/volcengine/OpenViking.git
synced 2026-10-01 09:48:03 +08:00
Qin Haojie
9295a3b955
Keep sessions user-scoped and remove legacy agent fallback that could leak an actor view into commit memory writes.
293 lines
10 KiB
Python
293 lines
10 KiB
Python
# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
|
|
# SPDX-License-Identifier: AGPL-3.0
|
|
|
|
"""Shared fixtures for OpenViking server tests."""
|
|
|
|
import asyncio
|
|
import shutil
|
|
import socket
|
|
import threading
|
|
import time
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
|
|
import httpx
|
|
import pytest
|
|
import pytest_asyncio
|
|
import uvicorn
|
|
|
|
from openviking import AsyncOpenViking
|
|
from openviking.models.embedder.base import DenseEmbedderBase, EmbedResult
|
|
from openviking.server.app import create_app
|
|
from openviking.server.config import ServerConfig
|
|
from openviking.server.identity import RequestContext, Role
|
|
from openviking.service.core import OpenVikingService
|
|
from openviking.storage.queuefs import QueueManager, SessionCommitMsg, get_queue_manager
|
|
from openviking_cli.session.user_id import UserIdentifier
|
|
from openviking_cli.utils.config.embedding_config import EmbeddingConfig
|
|
from openviking_cli.utils.config.vlm_config import VLMConfig
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Paths
|
|
# ---------------------------------------------------------------------------
|
|
|
|
PROJECT_ROOT = Path(__file__).parent.parent.parent
|
|
TEST_TMP_DIR = PROJECT_ROOT / "test_data" / "tmp_server"
|
|
SDK_ROOT_API_KEY = "test-root-api-key"
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Sample data
|
|
# ---------------------------------------------------------------------------
|
|
|
|
SAMPLE_MD_CONTENT = """\
|
|
# Sample Document
|
|
|
|
## Introduction
|
|
This is a sample markdown document for server testing.
|
|
|
|
## Features
|
|
- Feature 1: Resource management
|
|
- Feature 2: Semantic search
|
|
"""
|
|
|
|
|
|
def _install_fake_embedder(monkeypatch):
|
|
"""Use an in-process fake embedder so server tests never hit external APIs."""
|
|
|
|
class FakeEmbedder(DenseEmbedderBase):
|
|
def __init__(self, dimension: int = 2048):
|
|
super().__init__(model_name="test-fake-embedder")
|
|
self._dimension = dimension
|
|
|
|
def embed(self, text: str, is_query: bool = False) -> EmbedResult:
|
|
return EmbedResult(dense_vector=[0.1] * self._dimension)
|
|
|
|
def get_dimension(self) -> int:
|
|
return self._dimension
|
|
|
|
monkeypatch.setattr(EmbeddingConfig, "get_embedder", lambda self: FakeEmbedder(self.dimension))
|
|
return FakeEmbedder
|
|
|
|
|
|
def _install_fake_vlm(monkeypatch):
|
|
"""Use a fake VLM so server tests never hit external LLM APIs."""
|
|
|
|
async def _fake_get_completion(self, prompt, thinking=False):
|
|
if "extract user-private configuration items" in prompt:
|
|
return '{"values": {"api_key": "secret-xyz", "base_url": "https://example.com"}}'
|
|
return "# Test Summary\n\nFake summary for testing.\n\n## Details\nTest content."
|
|
|
|
async def _fake_get_vision_completion(self, prompt, images, thinking=False):
|
|
return "Fake image description for testing."
|
|
|
|
monkeypatch.setattr(VLMConfig, "is_available", lambda self: True)
|
|
monkeypatch.setattr(VLMConfig, "get_completion_async", _fake_get_completion)
|
|
monkeypatch.setattr(VLMConfig, "get_vision_completion_async", _fake_get_vision_completion)
|
|
|
|
|
|
def _install_session_commit_queue_fallback(service: OpenVikingService, monkeypatch) -> None:
|
|
"""Execute SessionCommit jobs when MockLocalAGFS cannot dequeue QueueFS entries."""
|
|
queue_manager = get_queue_manager()
|
|
original_enqueue = queue_manager.enqueue
|
|
|
|
async def enqueue_with_session_commit_fallback(queue_name, data):
|
|
if queue_name != QueueManager.SESSION_COMMIT:
|
|
return await original_enqueue(queue_name, data)
|
|
|
|
async def process_commit() -> None:
|
|
await asyncio.sleep(0)
|
|
msg = SessionCommitMsg(**data)
|
|
ctx = RequestContext(
|
|
user=UserIdentifier.from_dict(msg.user),
|
|
role=Role.USER,
|
|
)
|
|
queued_session = service.sessions.session(
|
|
ctx,
|
|
msg.session_id,
|
|
session_uri=msg.session_uri,
|
|
)
|
|
await queued_session.load()
|
|
await queued_session.resume_queued_commit(msg)
|
|
|
|
asyncio.create_task(process_commit())
|
|
return data["task_id"]
|
|
|
|
monkeypatch.setattr(queue_manager, "enqueue", enqueue_with_session_commit_fallback)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Core fixtures: service + app + async client (HTTP API tests, in-process)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def temp_dir():
|
|
"""Create a unique temp directory per test, auto-cleanup."""
|
|
import uuid
|
|
|
|
unique_dir = TEST_TMP_DIR / uuid.uuid4().hex[:8]
|
|
unique_dir.mkdir(parents=True, exist_ok=True)
|
|
yield unique_dir
|
|
shutil.rmtree(unique_dir, ignore_errors=True)
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def sample_markdown_file(temp_dir: Path) -> Path:
|
|
"""Create a sample markdown file for resource tests."""
|
|
f = temp_dir / "sample.md"
|
|
f.write_text(SAMPLE_MD_CONTENT)
|
|
return f
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def upload_temp_dir(temp_dir: Path, monkeypatch) -> Path:
|
|
"""Use the per-test temp directory as the HTTP upload temp dir."""
|
|
config = SimpleNamespace(
|
|
storage=SimpleNamespace(get_upload_temp_dir=lambda: temp_dir),
|
|
)
|
|
monkeypatch.setattr(
|
|
"openviking.server.routers.resources.get_openviking_config",
|
|
lambda: config,
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
"openviking.server.routers.pack.get_openviking_config",
|
|
lambda: config,
|
|
raising=False,
|
|
)
|
|
monkeypatch.setattr(
|
|
"openviking.server.temp_upload_store.get_openviking_config",
|
|
lambda: config,
|
|
)
|
|
return temp_dir
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def service(temp_dir: Path, monkeypatch):
|
|
"""Create and initialize an OpenVikingService in embedded mode."""
|
|
fake_embedder_cls = _install_fake_embedder(monkeypatch)
|
|
_install_fake_vlm(monkeypatch)
|
|
svc = OpenVikingService(
|
|
path=str(temp_dir / "data"), user=UserIdentifier.the_default_user("test_user")
|
|
)
|
|
await svc.initialize()
|
|
svc.viking_fs.query_embedder = fake_embedder_cls()
|
|
_install_session_commit_queue_fallback(svc, monkeypatch)
|
|
yield svc
|
|
await svc.close()
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def app(service: OpenVikingService):
|
|
"""Create FastAPI app with pre-initialized service (no auth)."""
|
|
from openviking.server.auth.plugins import DevAuthPlugin
|
|
from openviking.server.auth.registry import get_registry
|
|
from openviking.server.dependencies import set_service
|
|
|
|
config = ServerConfig()
|
|
fastapi_app = create_app(config=config, service=service)
|
|
# ASGITransport doesn't trigger lifespan, so wire up the service manually
|
|
set_service(service)
|
|
# Manually initialize auth plugin (lifespan not triggered in ASGI tests)
|
|
registry = get_registry()
|
|
if registry.get("dev") is None:
|
|
registry.register(DevAuthPlugin)
|
|
plugin_cls = registry.get("dev")
|
|
if plugin_cls is not None:
|
|
fastapi_app.state.auth_plugin = plugin_cls()
|
|
return fastapi_app
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def client(app):
|
|
"""httpx AsyncClient bound to the ASGI app (no real network)."""
|
|
transport = httpx.ASGITransport(app=app)
|
|
async with httpx.AsyncClient(transport=transport, base_url="http://testserver") as c:
|
|
yield c
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def client_with_resource(client, service, sample_markdown_file):
|
|
"""Client + a resource already added and processed."""
|
|
ctx = RequestContext(user=UserIdentifier.the_default_user(), role=Role.ROOT)
|
|
result = await service.resources.add_resource(
|
|
path=str(sample_markdown_file),
|
|
ctx=ctx,
|
|
reason="test resource",
|
|
wait=True,
|
|
)
|
|
yield client, result.get("root_uri", "")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# SDK fixtures: real uvicorn server + AsyncHTTPClient (end-to-end tests)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest_asyncio.fixture(scope="function")
|
|
async def running_server(temp_dir: Path, monkeypatch):
|
|
"""Start a real uvicorn server in a background thread."""
|
|
await AsyncOpenViking.reset()
|
|
fake_embedder_cls = _install_fake_embedder(monkeypatch)
|
|
_install_fake_vlm(monkeypatch)
|
|
|
|
@asynccontextmanager
|
|
async def _noop_mcp_lifespan():
|
|
yield
|
|
|
|
monkeypatch.setattr("openviking.server.mcp_endpoint.mcp_lifespan", _noop_mcp_lifespan)
|
|
|
|
svc = OpenVikingService(
|
|
path=str(temp_dir / "sdk_data"), user=UserIdentifier.the_default_user("sdk_test_user")
|
|
)
|
|
await svc.initialize()
|
|
svc.viking_fs.query_embedder = fake_embedder_cls()
|
|
_install_session_commit_queue_fallback(svc, monkeypatch)
|
|
|
|
config = ServerConfig(root_api_key=SDK_ROOT_API_KEY)
|
|
fastapi_app = create_app(config=config, service=svc)
|
|
|
|
# Find a free port
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
port = s.getsockname()[1]
|
|
|
|
uvi_config = uvicorn.Config(fastapi_app, host="127.0.0.1", port=port, log_level="warning")
|
|
server = uvicorn.Server(uvi_config)
|
|
thread = threading.Thread(target=server.run, daemon=True)
|
|
thread.start()
|
|
|
|
# Wait for server ready
|
|
for _ in range(50):
|
|
try:
|
|
r = httpx.get(f"http://127.0.0.1:{port}/health", timeout=1)
|
|
if r.status_code == 200:
|
|
break
|
|
except Exception:
|
|
time.sleep(0.1)
|
|
|
|
for _ in range(50):
|
|
if getattr(fastapi_app.state, "api_key_manager", None) is not None:
|
|
break
|
|
time.sleep(0.1)
|
|
else:
|
|
raise RuntimeError("APIKeyManager did not initialize for SDK server test")
|
|
|
|
manager = fastapi_app.state.api_key_manager
|
|
sdk_account_id = "sdk_test_account"
|
|
sdk_user_key = await manager.create_account(sdk_account_id, "sdk_test_user")
|
|
sdk_ctx = RequestContext(
|
|
user=UserIdentifier(sdk_account_id, "sdk_test_user"),
|
|
role=Role.ADMIN,
|
|
)
|
|
await svc.initialize_account_directories(sdk_ctx)
|
|
await svc.initialize_user_directories(sdk_ctx)
|
|
|
|
yield port, svc, sdk_user_key
|
|
|
|
server.should_exit = True
|
|
thread.join(timeout=5)
|
|
await svc.close()
|
|
await AsyncOpenViking.reset()
|