deer-flow/backend/tests/blocking_io/test_managed_models.py

"""Managed model encryption, persistence and config merging stay off the loop."""

import asyncio
import json
from types import SimpleNamespace

import httpx
import pytest
from langchain_openai import ChatOpenAI

from app.gateway.routers import managed_models as router
from deerflow.config.app_config import AppConfig
from deerflow.config.managed_models import ManagedModel, ManagedModelStore
from deerflow.models.patched_deepseek import PatchedChatDeepSeek


@pytest.mark.asyncio
async def test_admin_catalog_round_trip_offloads_storage(tmp_path, monkeypatch):
    monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
    config = AppConfig.model_validate({"sandbox": {"use": "test"}})
    monkeypatch.setattr(router, "get_app_config", lambda: config)
    request = SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(system_role="admin")))
    body = router.SaveModelRequest(config=ManagedModel(name="test", model="test", base_url="https://example.com/v1", api_key="secret"))
    result = await router.save_model(request, body)
    assert result["has_api_key"] is True
    catalog = await router.list_managed_models(request)
    assert catalog["models"][0]["name"] == "test"
    assert "secret" not in str(catalog)


@pytest.mark.asyncio
async def test_deepseek_probe_offloads_credentials_and_client_construction(tmp_path, monkeypatch):
    monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
    monkeypatch.setenv("LANGSMITH_TRACING", "false")
    store = await asyncio.to_thread(ManagedModelStore)
    profile = ManagedModel(name="flash", model="deepseek-flash", base_url="https://api.deepseek.com", api_key="test-secret")
    saved = await asyncio.to_thread(store.save, profile, expected_revision=None)
    original_catalog = await asyncio.to_thread(store.path.read_bytes)
    constructed = []

    def observe_constructor(original):
        def initialize(self, *args, **kwargs):
            with pytest.raises(RuntimeError, match="no running event loop"):
                asyncio.get_running_loop()
            original(self, *args, **kwargs)
            constructed.append(type(self).__name__)

        return initialize

    monkeypatch.setattr(ChatOpenAI, "__init__", observe_constructor(ChatOpenAI.__init__))
    monkeypatch.setattr(PatchedChatDeepSeek, "__init__", observe_constructor(PatchedChatDeepSeek.__init__))

    async def send(client, request, **kwargs):
        assert request.headers["authorization"] == "Bearer test-secret"
        body = json.loads(request.content)
        assert body["thinking"] == {"type": "disabled"}
        chunk = {
            "id": "chat-test",
            "object": "chat.completion.chunk",
            "created": 1,
            "model": "deepseek-flash",
            "choices": [
                {"index": 0, "delta": {"role": "assistant", "content": "", "tool_calls": [{"index": 0, "id": "call-test", "type": "function", "function": {"name": "connection_check", "arguments": "{}"}}]}, "finish_reason": "tool_calls"}
            ],
        }
        return httpx.Response(200, request=request, headers={"content-type": "text/event-stream"}, content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n".encode())

    monkeypatch.setattr(httpx.AsyncClient, "send", send)
    request = SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(system_role="admin")))
    draft = ManagedModel(name="flash", model="deepseek-flash", base_url="https://api.deepseek.com")
    result = await router.test_model(request, router.SaveModelRequest(config=draft, expected_revision=saved.revision))
    assert result == {"ok": True, "message": "success"}
    assert constructed == ["PatchedChatDeepSeek"]
    assert await asyncio.to_thread(store.path.read_bytes) == original_catalog


@pytest.mark.asyncio
async def test_preset_catalog_registration_and_native_client_stay_off_loop(tmp_path, monkeypatch):
    from langchain_anthropic import ChatAnthropic

    monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path))
    monkeypatch.setenv("LANGSMITH_TRACING", "false")
    base = AppConfig.model_validate({"sandbox": {"use": "test"}})
    monkeypatch.setattr(router, "get_app_config", lambda: base)
    request = SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(system_role="admin")))
    providers = await router.list_model_providers(request)
    assert any(item["id"] == "anthropic" for item in providers["providers"])
    result = await router.save_model(request, router.SaveModelRequest(config=ManagedModel(provider="anthropic", model="arbitrary", api_key="test-secret")))
    assert result["name"] and result["provider"] == "anthropic"
    assert "test-secret" not in str(await router.list_managed_models(request))
    original = ChatAnthropic.__init__
    built = []

    def initialize(self, *args, **kwargs):
        with pytest.raises(RuntimeError, match="no running event loop"):
            asyncio.get_running_loop()
        original(self, *args, **kwargs)
        built.append(True)

    monkeypatch.setattr(ChatAnthropic, "__init__", initialize)

    # Stop at the wire boundary; constructing the actual native SDK is the
    # assertion here. The real streaming wire test lives in the offline suite.
    async def send(client, outbound, **kwargs):
        assert outbound.headers["x-api-key"] == "test-secret"
        return httpx.Response(401, request=outbound, json={"error": {"type": "authentication_error", "message": "private provider detail"}})

    monkeypatch.setattr(httpx.AsyncClient, "send", send)
    profile = ManagedModel(name=result["name"], provider="anthropic", model="arbitrary")
    assert await router.test_model(request, router.SaveModelRequest(config=profile, expected_revision=result["revision"])) == {"ok": False, "message": "authentication_failed"}
    assert built == [True]