OpenViking/tests/server/conftest.py

# 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.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.service.task_tracker import get_task_tracker
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
    queued_commit_tasks: set[str] = set()
    tracker = get_task_tracker()
    original_has_work = tracker.has_work
    monkeypatch.setattr(
        tracker,
        "has_work",
        lambda task_id: task_id in queued_commit_tasks or original_has_work(task_id),
    )

    async def enqueue_with_session_commit_fallback(queue_name, data):
        if queue_name != QueueManager.SESSION_COMMIT:
            return await original_enqueue(queue_name, data)

        queued_commit_tasks.add(data["task_id"])

        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,
            )
            while True:
                phase1 = await queued_session._read_phase1_meta(msg.archive_uri)
                if phase1.get("status") == "ready" or await queued_session._archive_file_exists(
                    msg.archive_uri,
                    ".failed.json",
                ):
                    break
                await asyncio.sleep(0)
            await queued_session.load()
            try:
                while not await queued_session.resume_queued_commit(msg):
                    await asyncio.sleep(0)
            finally:
                queued_commit_tasks.discard(msg.task_id)

        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 for in-process API tests."""
    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."""
    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()