chore: import upstream snapshot with attribution
This commit is contained in:
@@ -0,0 +1,183 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import asyncpg
|
||||
import pytest
|
||||
from yuxi.services.run_queue_service import append_run_stream_event, get_redis_client
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
def _postgres_dsn() -> str:
|
||||
return os.getenv("POSTGRES_URL", "postgresql+asyncpg://postgres:postgres@postgres:5432/yuxi").replace(
|
||||
"+asyncpg", ""
|
||||
)
|
||||
|
||||
|
||||
async def _collect_sse_payloads(response) -> list[tuple[str, dict, str | None]]:
|
||||
event = "message"
|
||||
event_id = None
|
||||
data_lines: list[str] = []
|
||||
payloads: list[tuple[str, dict, str | None]] = []
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line:
|
||||
if data_lines:
|
||||
payloads.append((event, json.loads("\n".join(data_lines)), event_id))
|
||||
if event == "end":
|
||||
return payloads
|
||||
event = "message"
|
||||
event_id = None
|
||||
data_lines = []
|
||||
continue
|
||||
if line.startswith(":"):
|
||||
continue
|
||||
if line.startswith("event:"):
|
||||
event = line.removeprefix("event:").strip() or "message"
|
||||
elif line.startswith("id:"):
|
||||
event_id = line.removeprefix("id:").strip()
|
||||
elif line.startswith("data:"):
|
||||
data_lines.append(line.removeprefix("data:").strip())
|
||||
|
||||
return payloads
|
||||
|
||||
|
||||
async def test_run_events_verbose_false_returns_compact_payload(test_client, standard_user):
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
run_id = str(uuid.uuid4())
|
||||
thread_id = str(uuid.uuid4())
|
||||
request_id = f"req-{uuid.uuid4()}"
|
||||
|
||||
conn = await asyncpg.connect(_postgres_dsn())
|
||||
try:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO agent_runs
|
||||
(id, conversation_thread_id, agent_slug, uid, request_id, input_payload, status, run_type)
|
||||
VALUES ($1, $2, $3, $4, $5, $6::jsonb, $7, $8)
|
||||
""",
|
||||
run_id,
|
||||
thread_id,
|
||||
"deep-research",
|
||||
uid,
|
||||
request_id,
|
||||
json.dumps({"query": "写一个冒泡排序"}, ensure_ascii=False),
|
||||
"completed",
|
||||
"chat",
|
||||
)
|
||||
finally:
|
||||
await conn.close()
|
||||
|
||||
try:
|
||||
await append_run_stream_event(
|
||||
run_id,
|
||||
"metadata",
|
||||
{
|
||||
"request_id": request_id,
|
||||
"agent_slug": "deep-research",
|
||||
"backend_id": "ChatbotAgent",
|
||||
"uid": uid,
|
||||
},
|
||||
thread_id=thread_id,
|
||||
)
|
||||
await append_run_stream_event(
|
||||
run_id,
|
||||
"custom",
|
||||
{
|
||||
"name": "yuxi.agent_state",
|
||||
"chunk": {
|
||||
"request_id": request_id,
|
||||
"response": None,
|
||||
"thread_id": thread_id,
|
||||
"status": "agent_state",
|
||||
"agent_state": {
|
||||
"todos": [],
|
||||
"files": {},
|
||||
"artifacts": [],
|
||||
"subagent_runs": [],
|
||||
},
|
||||
"meta": {"uid": uid},
|
||||
},
|
||||
"agent_state": {
|
||||
"todos": [],
|
||||
"files": {},
|
||||
"artifacts": [],
|
||||
"subagent_runs": [],
|
||||
},
|
||||
},
|
||||
thread_id=thread_id,
|
||||
)
|
||||
await append_run_stream_event(
|
||||
run_id,
|
||||
"messages",
|
||||
{
|
||||
"items": [
|
||||
{
|
||||
"request_id": request_id,
|
||||
"response": "你",
|
||||
"thread_id": thread_id,
|
||||
"status": "loading",
|
||||
"stream_event": {
|
||||
"type": "tool_call",
|
||||
"message_id": "msg-1",
|
||||
"tool_call_id": "call-1",
|
||||
"name": "ls",
|
||||
"args": {"path": "/home/gem/user-data/outputs"},
|
||||
"thread_id": thread_id,
|
||||
"namespace": [],
|
||||
},
|
||||
"metadata": {
|
||||
"langfuse_user_id": uid,
|
||||
"langgraph_checkpoint_ns": "model:checkpoint",
|
||||
},
|
||||
}
|
||||
]
|
||||
},
|
||||
thread_id=thread_id,
|
||||
)
|
||||
await append_run_stream_event(
|
||||
run_id,
|
||||
"end",
|
||||
{"status": "completed", "chunk": {"status": "finished", "request_id": request_id, "meta": {"uid": uid}}},
|
||||
thread_id=thread_id,
|
||||
)
|
||||
|
||||
async with test_client.stream(
|
||||
"GET",
|
||||
f"/api/agent/runs/{run_id}/events",
|
||||
params={"verbose": "false"},
|
||||
headers=standard_user["headers"],
|
||||
) as response:
|
||||
assert response.status_code == 200, response.text
|
||||
payloads = await _collect_sse_payloads(response)
|
||||
|
||||
assert {event for event, _payload, _event_id in payloads} == {"messages", "end"}
|
||||
|
||||
message_event = next(item for item in payloads if item[0] == "messages")
|
||||
message_chunk = message_event[1]["payload"]["items"][0]
|
||||
assert message_event[1]["request_id"] == request_id
|
||||
assert message_event[2]
|
||||
assert "request_id" not in message_chunk
|
||||
assert "metadata" not in message_chunk
|
||||
assert "response" not in message_chunk
|
||||
assert "thread_id" not in message_chunk
|
||||
assert message_chunk["stream_event"]["tool_call_id"] == "call-1"
|
||||
assert "thread_id" not in message_chunk["stream_event"]
|
||||
assert "namespace" not in message_chunk["stream_event"]
|
||||
|
||||
end_event = next(item for item in payloads if item[0] == "end")
|
||||
assert end_event[1]["request_id"] == request_id
|
||||
assert end_event[1]["payload"]["status"] == "completed"
|
||||
assert "request_id" not in end_event[1]["payload"]["chunk"]
|
||||
assert "meta" not in end_event[1]["payload"]["chunk"]
|
||||
finally:
|
||||
redis = await get_redis_client()
|
||||
await redis.delete(f"run:events:{run_id}")
|
||||
conn = await asyncpg.connect(_postgres_dsn())
|
||||
try:
|
||||
await conn.execute("DELETE FROM agent_runs WHERE id = $1", run_id)
|
||||
finally:
|
||||
await conn.close()
|
||||
@@ -0,0 +1,190 @@
|
||||
"""
|
||||
Integration tests for API Key router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
API_KEYS_PATH = "/api/user/apikey/"
|
||||
# 受登录保护的轻量端点,用于验证 Bearer 鉴权(API Key / JWT)是否生效,无需执行智能体
|
||||
PROTECTED_PATH = "/api/agent"
|
||||
|
||||
|
||||
async def test_list_api_keys_requires_auth(test_client):
|
||||
"""List API keys should require authentication."""
|
||||
response = await test_client.get(API_KEYS_PATH)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_list_api_keys_requires_admin(test_client, admin_headers):
|
||||
"""List API keys should require admin privileges."""
|
||||
response = await test_client.get(API_KEYS_PATH, headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert "api_keys" in data
|
||||
assert "total" in data
|
||||
|
||||
|
||||
async def test_create_api_key(test_client, admin_headers):
|
||||
"""Admin should be able to create a new API key."""
|
||||
payload = {
|
||||
"name": "Test API Key",
|
||||
}
|
||||
response = await test_client.post(API_KEYS_PATH, json=payload, headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert "api_key" in data
|
||||
assert "secret" in data
|
||||
assert data["api_key"]["name"] == "Test API Key"
|
||||
assert data["api_key"]["key_prefix"].startswith("yxkey_")
|
||||
assert data["secret"].startswith(data["api_key"]["key_prefix"])
|
||||
|
||||
|
||||
async def test_get_api_key(test_client, admin_headers):
|
||||
"""Admin should be able to get a single API key."""
|
||||
# First create a key
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Get Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
# Then retrieve it
|
||||
response = await test_client.get(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert data["api_key"]["id"] == created["id"]
|
||||
assert data["api_key"]["name"] == "Get Test"
|
||||
|
||||
|
||||
async def test_update_api_key(test_client, admin_headers):
|
||||
"""Admin should be able to update an API key."""
|
||||
# Create a key
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Update Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
# Update it
|
||||
response = await test_client.put(
|
||||
f"{API_KEYS_PATH}{created['id']}",
|
||||
json={"name": "Updated Name", "is_enabled": False},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert data["api_key"]["name"] == "Updated Name"
|
||||
assert data["api_key"]["is_enabled"] is False
|
||||
|
||||
|
||||
async def test_delete_api_key(test_client, admin_headers):
|
||||
"""Admin should be able to delete an API key."""
|
||||
# Create a key
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Delete Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
# Delete it
|
||||
response = await test_client.delete(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["success"] is True
|
||||
|
||||
# Verify it's gone
|
||||
get_response = await test_client.get(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
assert get_response.status_code == 404
|
||||
|
||||
|
||||
async def test_regenerate_api_key(test_client, admin_headers):
|
||||
"""Admin should be able to regenerate an API key."""
|
||||
# Create a key
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Regenerate Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
original_secret = create_response.json()["secret"]
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
# Regenerate it
|
||||
response = await test_client.post(f"{API_KEYS_PATH}{created['id']}/regenerate", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
data = response.json()
|
||||
assert "secret" in data
|
||||
assert data["secret"] != original_secret
|
||||
assert data["api_key"]["key_prefix"] != original_secret[:12]
|
||||
|
||||
|
||||
async def test_api_key_auth_protected_endpoint(test_client, admin_headers):
|
||||
"""Test that API Key can be used to authenticate to a protected endpoint via Bearer token."""
|
||||
# Create an API key
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Auth Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
api_key_secret = create_response.json()["secret"]
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
try:
|
||||
response = await test_client.get(
|
||||
PROTECTED_PATH,
|
||||
headers={"Authorization": f"Bearer {api_key_secret}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
finally:
|
||||
# Cleanup: delete the test API key
|
||||
await test_client.delete(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
|
||||
|
||||
async def test_api_key_auth_requires_valid_key(test_client):
|
||||
"""Test that invalid API Key is rejected."""
|
||||
# Call protected endpoint with invalid API Key
|
||||
response = await test_client.get(
|
||||
PROTECTED_PATH,
|
||||
headers={"Authorization": "Bearer yxkey_invalid_key_that_does_not_exist"},
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
|
||||
|
||||
async def test_api_key_auth_requires_bearer_prefix(test_client, admin_headers):
|
||||
"""Test that API Key must be prefixed with 'Bearer '."""
|
||||
# Create an API key
|
||||
admin_response = await test_client.post(API_KEYS_PATH, json={"name": "Prefix Test"}, headers=admin_headers)
|
||||
assert admin_response.status_code == 200
|
||||
api_key_secret = admin_response.json()["secret"]
|
||||
created = admin_response.json()["api_key"]
|
||||
|
||||
try:
|
||||
# Call without Bearer prefix should fail
|
||||
response = await test_client.get(
|
||||
PROTECTED_PATH,
|
||||
headers={"Authorization": api_key_secret}, # Missing "Bearer " prefix
|
||||
)
|
||||
assert response.status_code == 401, response.text
|
||||
finally:
|
||||
# Cleanup: delete the test API key
|
||||
await test_client.delete(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
|
||||
|
||||
async def test_jwt_still_works_after_apikey_auth(test_client, admin_headers):
|
||||
"""Test that JWT Bearer tokens still work after API Key changes."""
|
||||
# Call protected endpoint with JWT Bearer token (admin_headers)
|
||||
response = await test_client.get(PROTECTED_PATH, headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
async def test_api_key_auto_binds_to_current_user(test_client, admin_headers):
|
||||
"""Test that API Key created without user_id is auto-bound to creator."""
|
||||
# Create API key as admin
|
||||
create_response = await test_client.post(API_KEYS_PATH, json={"name": "Auto Bind Test"}, headers=admin_headers)
|
||||
assert create_response.status_code == 200
|
||||
created = create_response.json()["api_key"]
|
||||
|
||||
try:
|
||||
# Verify user_id is set (auto-bound to admin)
|
||||
assert created["user_id"] is not None, "API Key should be auto-bound to creator"
|
||||
|
||||
# Verify the key can be used for auth
|
||||
api_key_secret = create_response.json()["secret"]
|
||||
response = await test_client.get(
|
||||
PROTECTED_PATH,
|
||||
headers={"Authorization": f"Bearer {api_key_secret}"},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
finally:
|
||||
# Cleanup: delete the test API key
|
||||
await test_client.delete(f"{API_KEYS_PATH}{created['id']}", headers=admin_headers)
|
||||
@@ -0,0 +1,254 @@
|
||||
"""
|
||||
Integration tests for authentication-related API routes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def _require_superadmin(test_client, headers):
|
||||
response = await test_client.get("/api/auth/me", headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
if response.json()["role"] != "superadmin":
|
||||
pytest.fail("This test requires TEST_USERNAME to be a superadmin account.")
|
||||
|
||||
|
||||
async def _create_department_with_admin(test_client, headers, label: str) -> dict:
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
admin_uid = f"adm{label}_{suffix}"
|
||||
admin_password = f"Pw!{suffix}"
|
||||
response = await test_client.post(
|
||||
"/api/departments",
|
||||
json={
|
||||
"name": f"pytest_{label}_{suffix}",
|
||||
"description": "pytest managed department",
|
||||
"admin_uid": admin_uid,
|
||||
"admin_password": admin_password,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 201, response.text
|
||||
|
||||
login_response = await test_client.post(
|
||||
"/api/auth/token",
|
||||
data={"username": admin_uid, "password": admin_password},
|
||||
)
|
||||
assert login_response.status_code == 200, login_response.text
|
||||
|
||||
login_payload = login_response.json()
|
||||
return {
|
||||
"department": response.json(),
|
||||
"admin_id": login_payload["user_id"],
|
||||
"admin_headers": {"Authorization": f"Bearer {login_payload['access_token']}"},
|
||||
}
|
||||
|
||||
|
||||
async def _create_user(test_client, headers, label: str, role: str = "user", department_id: int | None = None) -> dict:
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
payload = {
|
||||
"username": f"u{label}_{suffix}",
|
||||
"password": f"Pw!{suffix}",
|
||||
"role": role,
|
||||
}
|
||||
if department_id is not None:
|
||||
payload["department_id"] = department_id
|
||||
|
||||
response = await test_client.post("/api/auth/users", json=payload, headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()
|
||||
|
||||
|
||||
async def _cleanup_user(test_client, headers, user_id: int) -> None:
|
||||
response = await test_client.delete(f"/api/auth/users/{user_id}", headers=headers)
|
||||
assert response.status_code in {200, 404}, response.text
|
||||
|
||||
|
||||
async def _cleanup_department(test_client, headers, department_id: int) -> None:
|
||||
response = await test_client.delete(f"/api/departments/{department_id}", headers=headers)
|
||||
assert response.status_code in {200, 404}, response.text
|
||||
|
||||
|
||||
async def test_login_with_invalid_credentials(test_client):
|
||||
response = await test_client.post("/api/auth/token", data={"username": "invalid", "password": "invalid"})
|
||||
assert response.status_code == 401
|
||||
assert "detail" in response.json()
|
||||
|
||||
|
||||
async def test_user_is_locked_after_repeated_failed_logins(test_client, standard_user):
|
||||
uid = standard_user["user"]["uid"]
|
||||
|
||||
for attempt in range(1, 5):
|
||||
response = await test_client.post("/api/auth/token", data={"username": uid, "password": "wrong-password"})
|
||||
assert response.status_code == 401, response.text
|
||||
assert response.json()["detail"] == "用户名或密码错误"
|
||||
|
||||
locked_response = await test_client.post("/api/auth/token", data={"username": uid, "password": "wrong-password"})
|
||||
assert locked_response.status_code == 423, locked_response.text
|
||||
assert "X-Lock-Remaining" in locked_response.headers
|
||||
assert "账户已被锁定" in locked_response.json()["detail"]
|
||||
|
||||
still_locked_response = await test_client.post(
|
||||
"/api/auth/token",
|
||||
data={"username": uid, "password": standard_user["password"]},
|
||||
)
|
||||
assert still_locked_response.status_code == 423, still_locked_response.text
|
||||
assert "X-Lock-Remaining" in still_locked_response.headers
|
||||
assert "登录被锁定" in still_locked_response.json()["detail"]
|
||||
|
||||
|
||||
async def test_admin_can_login_and_fetch_profile(test_client, admin_headers):
|
||||
profile_response = await test_client.get("/api/auth/me", headers=admin_headers)
|
||||
assert profile_response.status_code == 200
|
||||
data = profile_response.json()
|
||||
assert data["role"] in {"admin", "superadmin"}
|
||||
assert data["username"]
|
||||
assert data["id"]
|
||||
|
||||
|
||||
async def test_profile_requires_authentication(test_client):
|
||||
response = await test_client.get("/api/auth/me")
|
||||
assert response.status_code == 401
|
||||
assert response.json()["detail"] == "请登录后再访问"
|
||||
|
||||
|
||||
async def test_admin_can_create_and_delete_user(test_client, admin_headers):
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
payload = {
|
||||
"username": f"rtu_{suffix}",
|
||||
"password": "routerTest123!",
|
||||
"role": "user",
|
||||
}
|
||||
create_response = await test_client.post("/api/auth/users", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
|
||||
created_user = create_response.json()
|
||||
assert created_user["username"] == payload["username"]
|
||||
assert created_user["role"] == payload["role"]
|
||||
|
||||
delete_response = await test_client.delete(f"/api/auth/users/{created_user['id']}", headers=admin_headers)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
delete_payload = delete_response.json()
|
||||
assert delete_payload["success"] is True
|
||||
assert delete_payload["message"] == "用户已删除"
|
||||
|
||||
|
||||
async def test_department_admin_is_limited_to_own_department_users(test_client, admin_headers):
|
||||
await _require_superadmin(test_client, admin_headers)
|
||||
|
||||
user_ids: list[int] = []
|
||||
admin_ids: list[int] = []
|
||||
department_ids: list[int] = []
|
||||
|
||||
try:
|
||||
dept_a = await _create_department_with_admin(test_client, admin_headers, "a")
|
||||
dept_b = await _create_department_with_admin(test_client, admin_headers, "b")
|
||||
department_a = dept_a["department"]
|
||||
department_b = dept_b["department"]
|
||||
department_ids.extend([department_a["id"], department_b["id"]])
|
||||
admin_ids.extend([dept_a["admin_id"], dept_b["admin_id"]])
|
||||
|
||||
user_a = await _create_user(test_client, dept_a["admin_headers"], "a")
|
||||
user_b = await _create_user(test_client, dept_b["admin_headers"], "b")
|
||||
superadmin_created_user = await _create_user(test_client, admin_headers, "s", department_id=department_b["id"])
|
||||
user_ids.extend([user_a["id"], user_b["id"], superadmin_created_user["id"]])
|
||||
|
||||
assert user_a["department_id"] == department_a["id"]
|
||||
assert superadmin_created_user["department_id"] == department_b["id"]
|
||||
|
||||
forbidden_create = await test_client.post(
|
||||
"/api/auth/users",
|
||||
json={
|
||||
"username": f"ux_{uuid.uuid4().hex[:8]}",
|
||||
"password": "routerTest123!",
|
||||
"role": "user",
|
||||
"department_id": department_b["id"],
|
||||
},
|
||||
headers=dept_a["admin_headers"],
|
||||
)
|
||||
assert forbidden_create.status_code == 403, forbidden_create.text
|
||||
|
||||
list_response = await test_client.get("/api/auth/users", headers=dept_a["admin_headers"])
|
||||
assert list_response.status_code == 200, list_response.text
|
||||
listed_users = list_response.json()
|
||||
listed_user_ids = {user["id"] for user in listed_users}
|
||||
assert user_a["id"] in listed_user_ids
|
||||
assert user_b["id"] not in listed_user_ids
|
||||
assert all(user["department_id"] == department_a["id"] for user in listed_users)
|
||||
|
||||
options_response = await test_client.get("/api/auth/users/access-options", headers=dept_a["admin_headers"])
|
||||
assert options_response.status_code == 200, options_response.text
|
||||
access_options = options_response.json()
|
||||
option_uids = {user["uid"] for user in access_options}
|
||||
assert user_a["uid"] in option_uids
|
||||
assert user_b["uid"] not in option_uids
|
||||
assert all(user["department_id"] == department_a["id"] for user in access_options)
|
||||
|
||||
superadmin_list_response = await test_client.get("/api/auth/users?limit=1000", headers=admin_headers)
|
||||
assert superadmin_list_response.status_code == 200, superadmin_list_response.text
|
||||
superadmin_user_ids = {user["id"] for user in superadmin_list_response.json()}
|
||||
assert user_a["id"] in superadmin_user_ids
|
||||
assert user_b["id"] in superadmin_user_ids
|
||||
|
||||
own_read = await test_client.get(f"/api/auth/users/{user_a['id']}", headers=dept_a["admin_headers"])
|
||||
assert own_read.status_code == 200, own_read.text
|
||||
|
||||
cross_read = await test_client.get(f"/api/auth/users/{user_b['id']}", headers=dept_a["admin_headers"])
|
||||
assert cross_read.status_code == 403, cross_read.text
|
||||
|
||||
cross_update = await test_client.put(
|
||||
f"/api/auth/users/{user_b['id']}",
|
||||
json={"username": f"ub_{uuid.uuid4().hex[:8]}"},
|
||||
headers=dept_a["admin_headers"],
|
||||
)
|
||||
assert cross_update.status_code == 403, cross_update.text
|
||||
|
||||
role_escalation = await test_client.put(
|
||||
f"/api/auth/users/{user_a['id']}", json={"role": "admin"}, headers=dept_a["admin_headers"]
|
||||
)
|
||||
assert role_escalation.status_code == 403, role_escalation.text
|
||||
|
||||
cross_delete = await test_client.delete(f"/api/auth/users/{user_b['id']}", headers=dept_a["admin_headers"])
|
||||
assert cross_delete.status_code == 403, cross_delete.text
|
||||
|
||||
own_delete = await test_client.delete(f"/api/auth/users/{user_a['id']}", headers=dept_a["admin_headers"])
|
||||
assert own_delete.status_code == 200, own_delete.text
|
||||
user_ids.remove(user_a["id"])
|
||||
finally:
|
||||
for user_id in user_ids:
|
||||
await _cleanup_user(test_client, admin_headers, user_id)
|
||||
for admin_id in admin_ids:
|
||||
await _cleanup_user(test_client, admin_headers, admin_id)
|
||||
for department_id in department_ids:
|
||||
await _cleanup_department(test_client, admin_headers, department_id)
|
||||
|
||||
|
||||
async def test_invalid_token_is_rejected(test_client):
|
||||
headers = {"Authorization": "Bearer not-a-real-token"}
|
||||
response = await test_client.get("/api/auth/me", headers=headers)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_deleted_user_token_is_rejected(test_client, admin_headers, standard_user):
|
||||
user_id = standard_user["user"]["id"]
|
||||
|
||||
delete_response = await test_client.delete(f"/api/auth/users/{user_id}", headers=admin_headers)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
|
||||
profile_response = await test_client.get("/api/auth/me", headers=standard_user["headers"])
|
||||
assert profile_response.status_code == 401
|
||||
|
||||
|
||||
async def test_locked_user_token_is_rejected(test_client, standard_user):
|
||||
uid = standard_user["user"]["uid"]
|
||||
|
||||
for _ in range(5):
|
||||
await test_client.post("/api/auth/token", data={"username": uid, "password": "wrong-password"})
|
||||
|
||||
profile_response = await test_client.get("/api/auth/me", headers=standard_user["headers"])
|
||||
assert profile_response.status_code == 423
|
||||
assert "X-Lock-Remaining" in profile_response.headers
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
Integration tests for current agent run endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_agent_run_endpoints_require_authentication(test_client):
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
create_response = await test_client.post(
|
||||
"/api/agent/runs",
|
||||
json={"query": "hello", "agent_slug": "default-chatbot", "thread_id": str(uuid.uuid4())},
|
||||
)
|
||||
assert create_response.status_code == 401
|
||||
assert (await test_client.get(f"/api/agent/runs/{run_id}")).status_code == 401
|
||||
assert (await test_client.post(f"/api/agent/runs/{run_id}/cancel")).status_code == 401
|
||||
|
||||
|
||||
async def test_agent_run_create_rejects_empty_input(test_client, admin_headers):
|
||||
response = await test_client.post(
|
||||
"/api/agent/runs",
|
||||
json={"query": "", "agent_slug": "default-chatbot", "thread_id": str(uuid.uuid4())},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 422
|
||||
|
||||
|
||||
async def test_agent_run_missing_resource_returns_not_found(test_client, admin_headers):
|
||||
run_id = str(uuid.uuid4())
|
||||
|
||||
get_response = await test_client.get(f"/api/agent/runs/{run_id}", headers=admin_headers)
|
||||
assert get_response.status_code == 404
|
||||
|
||||
cancel_response = await test_client.post(f"/api/agent/runs/{run_id}/cancel", headers=admin_headers)
|
||||
assert cancel_response.status_code == 404
|
||||
@@ -0,0 +1,269 @@
|
||||
"""
|
||||
Integration tests for chat router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from yuxi.agents.backends.sandbox import ensure_thread_dirs, sandbox_user_data_dir, sandbox_workspace_dir
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_chat_endpoints_require_authentication(test_client):
|
||||
assert (await test_client.get("/api/chat/threads")).status_code == 401
|
||||
assert (await test_client.get("/api/agent")).status_code == 401
|
||||
|
||||
|
||||
async def test_image_upload_composites_transparent_png_pixels_on_white(test_client, admin_headers):
|
||||
image = Image.new("RGBA", (2, 2), (255, 255, 255, 0))
|
||||
image.putpixel((0, 0), (50, 87, 244, 0))
|
||||
image.putpixel((1, 0), (50, 87, 244, 255))
|
||||
|
||||
with io.BytesIO() as buffer:
|
||||
image.save(buffer, format="PNG")
|
||||
image_bytes = buffer.getvalue()
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/chat/image/upload",
|
||||
headers=admin_headers,
|
||||
files={"file": ("transparent.png", image_bytes, "image/png")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert payload["mime_type"] == "image/png"
|
||||
|
||||
processed_data = base64.b64decode(payload["image_content"])
|
||||
with Image.open(io.BytesIO(processed_data)) as processed_image:
|
||||
rgb_image = processed_image.convert("RGB")
|
||||
|
||||
assert rgb_image.getpixel((0, 0)) == (255, 255, 255)
|
||||
assert rgb_image.getpixel((1, 0)) == (50, 87, 244)
|
||||
|
||||
|
||||
async def test_thread_artifact_uses_image_signature_for_content_type(test_client, admin_headers):
|
||||
thread_id = await _create_thread_for_user(test_client, admin_headers)
|
||||
image = Image.new("RGBA", (2, 2), (255, 255, 255, 0))
|
||||
|
||||
with io.BytesIO() as buffer:
|
||||
image.save(buffer, format="PNG")
|
||||
image_bytes = buffer.getvalue()
|
||||
|
||||
upload_response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/attachments",
|
||||
headers=admin_headers,
|
||||
files={"file": ("mislabeled.jpg", image_bytes, "image/jpeg")},
|
||||
)
|
||||
|
||||
assert upload_response.status_code == 200, upload_response.text
|
||||
attachment = upload_response.json()
|
||||
|
||||
artifact_response = await test_client.get(attachment["original_artifact_url"], headers=admin_headers)
|
||||
|
||||
assert artifact_response.status_code == 200, artifact_response.text
|
||||
assert artifact_response.headers["content-type"].startswith("image/png")
|
||||
assert artifact_response.content.startswith(b"\x89PNG\r\n\x1a\n")
|
||||
|
||||
|
||||
async def _create_thread_for_user(test_client, headers: dict[str, str]) -> str:
|
||||
agents_resp = await test_client.get("/api/agent", headers=headers)
|
||||
assert agents_resp.status_code == 200, agents_resp.text
|
||||
agents = agents_resp.json().get("agents", [])
|
||||
if not agents:
|
||||
pytest.skip("No agents available for chat router integration tests.")
|
||||
|
||||
agent_id = agents[0].get("agent_id") or agents[0].get("slug")
|
||||
if not agent_id:
|
||||
pytest.skip("Agent payload missing slug field.")
|
||||
|
||||
create_resp = await test_client.post(
|
||||
"/api/chat/thread",
|
||||
json={
|
||||
"agent_id": agent_id,
|
||||
"title": f"chat-router-test-{uuid.uuid4().hex[:8]}",
|
||||
"metadata": {},
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert create_resp.status_code == 200, create_resp.text
|
||||
payload = create_resp.json()
|
||||
thread_id = payload.get("thread_id") or payload.get("id")
|
||||
assert thread_id, f"Create thread response missing thread identifier: {payload}"
|
||||
return thread_id
|
||||
|
||||
|
||||
async def test_admin_can_list_agents(test_client, admin_headers):
|
||||
response = await test_client.get("/api/agent", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert isinstance(payload["agents"], list)
|
||||
if payload["agents"]:
|
||||
assert "agent_id" in payload["agents"][0]
|
||||
|
||||
|
||||
async def test_admin_can_read_default_agent(test_client, admin_headers):
|
||||
response = await test_client.get("/api/agent/default", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
agent = response.json()["agent"]
|
||||
assert agent["is_default"] is True
|
||||
assert agent["agent_id"]
|
||||
|
||||
|
||||
async def test_agent_detail_filters_configurable_items_by_role(
|
||||
test_client,
|
||||
admin_headers,
|
||||
standard_user,
|
||||
):
|
||||
agents_response = await test_client.get("/api/agent", headers=standard_user["headers"])
|
||||
assert agents_response.status_code == 200, agents_response.text
|
||||
agents = agents_response.json().get("agents", [])
|
||||
if not agents:
|
||||
pytest.skip("No agents are registered in the system.")
|
||||
|
||||
agent_id = agents[0].get("agent_id") or agents[0].get("slug")
|
||||
if not agent_id:
|
||||
pytest.skip("Agent payload missing slug field.")
|
||||
|
||||
user_agent_response = await test_client.get(f"/api/agent/{agent_id}", headers=standard_user["headers"])
|
||||
assert user_agent_response.status_code == 200, user_agent_response.text
|
||||
user_items = user_agent_response.json()["agent"].get("configurable_items", {})
|
||||
assert "summary_threshold" not in user_items
|
||||
assert "summary_keep_messages" not in user_items
|
||||
assert "summary_prompt" not in user_items
|
||||
assert "summary_tool_result_token_limit" not in user_items
|
||||
assert "max_execution_steps" not in user_items
|
||||
|
||||
admin_agent_response = await test_client.get(f"/api/agent/{agent_id}", headers=admin_headers)
|
||||
assert admin_agent_response.status_code == 200, admin_agent_response.text
|
||||
admin_items = admin_agent_response.json()["agent"].get("configurable_items", {})
|
||||
assert "summary_threshold" in admin_items
|
||||
assert "summary_keep_messages" in admin_items
|
||||
assert "summary_prompt" in admin_items
|
||||
assert "summary_tool_result_token_limit" in admin_items
|
||||
assert "max_execution_steps" in admin_items
|
||||
|
||||
|
||||
async def test_setting_default_agent_requires_admin(test_client, admin_headers, standard_user):
|
||||
agents_response = await test_client.get("/api/agent", headers=admin_headers)
|
||||
assert agents_response.status_code == 200, agents_response.text
|
||||
agents = agents_response.json().get("agents", [])
|
||||
|
||||
if not agents:
|
||||
pytest.skip("No agents are registered in the system.")
|
||||
|
||||
candidate_agent_id = agents[0].get("agent_id") or agents[0].get("slug")
|
||||
if not candidate_agent_id:
|
||||
pytest.skip("Agent payload missing slug field.")
|
||||
|
||||
forbidden_response = await test_client.post(
|
||||
f"/api/agent/{candidate_agent_id}/set_default",
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
assert forbidden_response.status_code == 403
|
||||
|
||||
update_response = await test_client.post(
|
||||
f"/api/agent/{candidate_agent_id}/set_default",
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert update_response.status_code == 200, update_response.text
|
||||
agent = update_response.json()["agent"]
|
||||
assert agent["agent_id"] == candidate_agent_id
|
||||
assert agent["is_default"] is True
|
||||
|
||||
|
||||
async def test_save_thread_artifact_to_workspace_copies_output_file(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
filename = f"artifact-{uuid.uuid4().hex[:8]}.md"
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
source_path = sandbox_user_data_dir(thread_id) / "outputs" / filename
|
||||
source_path.write_text("# artifact\n", encoding="utf-8")
|
||||
|
||||
response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/artifacts/save",
|
||||
json={"path": f"/home/gem/user-data/outputs/{filename}"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
assert payload["name"] == filename
|
||||
assert payload["source_path"] == f"/home/gem/user-data/outputs/{filename}"
|
||||
assert payload["saved_path"] == f"/home/gem/user-data/workspace/saved_artifacts/{filename}"
|
||||
|
||||
saved_path = sandbox_workspace_dir(thread_id, uid) / "saved_artifacts" / filename
|
||||
assert saved_path.exists()
|
||||
assert saved_path.read_text(encoding="utf-8") == "# artifact\n"
|
||||
|
||||
download_response = await test_client.get(payload["saved_artifact_url"], headers=headers)
|
||||
assert download_response.status_code == 200, download_response.text
|
||||
assert download_response.text == "# artifact\n"
|
||||
|
||||
|
||||
async def test_save_thread_artifact_to_workspace_auto_renames_conflicts(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
filename = f"artifact-{uuid.uuid4().hex[:8]}.txt"
|
||||
renamed_filename = filename.replace(".txt", " (1).txt")
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
source_path = sandbox_user_data_dir(thread_id) / "outputs" / filename
|
||||
source_path.write_text("first\n", encoding="utf-8")
|
||||
|
||||
first_response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/artifacts/save",
|
||||
json={"path": f"/home/gem/user-data/outputs/{filename}"},
|
||||
headers=headers,
|
||||
)
|
||||
assert first_response.status_code == 200, first_response.text
|
||||
|
||||
source_path.write_text("second\n", encoding="utf-8")
|
||||
second_response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/artifacts/save",
|
||||
json={"path": f"/home/gem/user-data/outputs/{filename}"},
|
||||
headers=headers,
|
||||
)
|
||||
assert second_response.status_code == 200, second_response.text
|
||||
|
||||
first_payload = first_response.json()
|
||||
second_payload = second_response.json()
|
||||
assert first_payload["saved_path"] == f"/home/gem/user-data/workspace/saved_artifacts/{filename}"
|
||||
assert second_payload["saved_path"] == f"/home/gem/user-data/workspace/saved_artifacts/{renamed_filename}"
|
||||
|
||||
first_saved = sandbox_workspace_dir(thread_id, uid) / "saved_artifacts" / filename
|
||||
second_saved = sandbox_workspace_dir(thread_id, uid) / "saved_artifacts" / renamed_filename
|
||||
assert first_saved.read_text(encoding="utf-8") == "first\n"
|
||||
assert second_saved.read_text(encoding="utf-8") == "second\n"
|
||||
|
||||
|
||||
async def test_save_thread_artifact_to_workspace_rejects_invalid_paths(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
invalid_response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/artifacts/save",
|
||||
json={"path": "/home/gem/user-data/not-allowed/demo.txt"},
|
||||
headers=headers,
|
||||
)
|
||||
assert invalid_response.status_code == 404, invalid_response.text
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
directory_path = sandbox_workspace_dir(thread_id, uid) / "nested-dir"
|
||||
directory_path.mkdir(parents=True, exist_ok=True)
|
||||
directory_response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/artifacts/save",
|
||||
json={"path": "/home/gem/user-data/workspace/nested-dir"},
|
||||
headers=headers,
|
||||
)
|
||||
assert directory_response.status_code == 400, directory_response.text
|
||||
@@ -0,0 +1,57 @@
|
||||
"""
|
||||
Integration tests for dashboard router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_dashboard_requires_authentication(test_client):
|
||||
response = await test_client.get("/api/dashboard/conversations")
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_standard_user_is_forbidden(test_client, standard_user):
|
||||
response = await test_client.get("/api/dashboard/conversations", headers=standard_user["headers"])
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
async def test_admin_can_fetch_conversations(test_client, admin_headers):
|
||||
response = await test_client.get("/api/dashboard/conversations", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
assert isinstance(response.json(), list)
|
||||
|
||||
|
||||
async def test_admin_can_fetch_stats(test_client, admin_headers):
|
||||
"""Test that all stats endpoints return 200 and don't crash on DB queries."""
|
||||
|
||||
# Test call timeseries stats for all types
|
||||
types = ["models", "agents", "tokens", "tools"]
|
||||
for stats_type in types:
|
||||
response = await test_client.get(
|
||||
f"/api/dashboard/stats/calls/timeseries?type={stats_type}&time_range=14days", headers=admin_headers
|
||||
)
|
||||
assert response.status_code == 200, f"{stats_type} stats failed: {response.text}"
|
||||
data = response.json()
|
||||
assert "data" in data
|
||||
assert "categories" in data
|
||||
|
||||
# Test user activity stats
|
||||
response = await test_client.get("/api/dashboard/stats/users", headers=admin_headers)
|
||||
assert response.status_code == 200, f"user stats failed: {response.text}"
|
||||
assert "total_users" in response.json()
|
||||
|
||||
# Test tool call stats
|
||||
response = await test_client.get("/api/dashboard/stats/tools", headers=admin_headers)
|
||||
assert response.status_code == 200, f"tool stats failed: {response.text}"
|
||||
assert "total_calls" in response.json()
|
||||
|
||||
|
||||
async def test_admin_can_fetch_feedbacks(test_client, admin_headers):
|
||||
"""Test that feedback endpoint returns 200 and handles the User join correctly."""
|
||||
response = await test_client.get("/api/dashboard/feedbacks", headers=admin_headers)
|
||||
assert response.status_code == 200, f"feedbacks failed: {response.text}"
|
||||
assert isinstance(response.json(), list)
|
||||
@@ -0,0 +1,99 @@
|
||||
"""
|
||||
Integration tests for department management API routes.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_superadmin_can_delete_department_with_users(test_client, admin_headers):
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
department_payload = {
|
||||
"name": f"pytest_department_{suffix}",
|
||||
"description": "integration test department",
|
||||
"admin_uid": f"pta_{suffix}",
|
||||
"admin_password": "RouterDept123!",
|
||||
}
|
||||
user_payload = {
|
||||
"username": f"dept_user_{suffix}",
|
||||
"password": "RouterUser123!",
|
||||
"role": "user",
|
||||
}
|
||||
|
||||
department_id = None
|
||||
created_user_id = None
|
||||
department_admin_id = None
|
||||
|
||||
try:
|
||||
create_department_response = await test_client.post(
|
||||
"/api/departments",
|
||||
json=department_payload,
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert create_department_response.status_code == 201, create_department_response.text
|
||||
department_id = create_department_response.json()["id"]
|
||||
|
||||
create_user_response = await test_client.post(
|
||||
"/api/auth/users",
|
||||
json={**user_payload, "department_id": department_id},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert create_user_response.status_code == 200, create_user_response.text
|
||||
created_user_id = create_user_response.json()["id"]
|
||||
|
||||
list_users_response = await test_client.get("/api/auth/users", headers=admin_headers)
|
||||
assert list_users_response.status_code == 200, list_users_response.text
|
||||
users_before_delete = list_users_response.json()
|
||||
department_admin = next((user for user in users_before_delete if user["uid"] == department_payload["admin_uid"]), None)
|
||||
assert department_admin is not None
|
||||
department_admin_id = department_admin["id"]
|
||||
|
||||
delete_department_response = await test_client.delete(f"/api/departments/{department_id}", headers=admin_headers)
|
||||
assert delete_department_response.status_code == 200, delete_department_response.text
|
||||
assert delete_department_response.json()["success"] is True
|
||||
department_id = None
|
||||
|
||||
deleted_department_response = await test_client.get(
|
||||
f"/api/departments/{create_department_response.json()['id']}",
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert deleted_department_response.status_code == 404, deleted_department_response.text
|
||||
|
||||
list_users_after_delete_response = await test_client.get("/api/auth/users", headers=admin_headers)
|
||||
assert list_users_after_delete_response.status_code == 200, list_users_after_delete_response.text
|
||||
users_after_delete = list_users_after_delete_response.json()
|
||||
|
||||
# 删除部门后用户迁移到默认部门(代码中默认部门固定为 id=1,其名称可被用户修改)
|
||||
migrated_admin = next((user for user in users_after_delete if user["id"] == department_admin_id), None)
|
||||
assert migrated_admin is not None
|
||||
assert migrated_admin["department_id"] == 1
|
||||
|
||||
migrated_user = next((user for user in users_after_delete if user["id"] == created_user_id), None)
|
||||
assert migrated_user is not None
|
||||
assert migrated_user["department_id"] == 1
|
||||
finally:
|
||||
if department_admin_id is not None:
|
||||
await test_client.delete(f"/api/auth/users/{department_admin_id}", headers=admin_headers)
|
||||
if created_user_id is not None:
|
||||
await test_client.delete(f"/api/auth/users/{created_user_id}", headers=admin_headers)
|
||||
if department_id is not None:
|
||||
await test_client.delete(f"/api/departments/{department_id}", headers=admin_headers)
|
||||
|
||||
|
||||
async def test_superadmin_cannot_delete_default_department(test_client, admin_headers):
|
||||
departments_response = await test_client.get("/api/departments", headers=admin_headers)
|
||||
assert departments_response.status_code == 200, departments_response.text
|
||||
departments = departments_response.json()
|
||||
|
||||
# 默认部门在代码中固定为 id=1(受保护不可删除),其名称可被用户修改,故按 id 定位
|
||||
default_department = next((department for department in departments if department["id"] == 1), None)
|
||||
assert default_department is not None
|
||||
|
||||
delete_response = await test_client.delete(f"/api/departments/{default_department['id']}", headers=admin_headers)
|
||||
assert delete_response.status_code == 400, delete_response.text
|
||||
assert delete_response.json()["detail"] == "默认部门不允许删除"
|
||||
@@ -0,0 +1,61 @@
|
||||
"""
|
||||
Integration tests for evaluation router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def _upload_test_dataset(test_client, admin_headers: dict[str, str], kb_id: str) -> tuple[str, str]:
|
||||
dataset_name = f"pytest_dataset_{uuid.uuid4().hex[:8]}"
|
||||
line = '{"query":"什么是单元测试?","gold_answer":"用于验证代码行为的自动化测试"}\n'
|
||||
|
||||
response = await test_client.post(
|
||||
f"/api/evaluation/databases/{kb_id}/datasets/upload",
|
||||
data={"name": dataset_name, "description": "pytest dataset for download"},
|
||||
files={"file": ("pytest_dataset.jsonl", line.encode("utf-8"), "application/x-ndjson")},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
assert payload.get("message") == "success"
|
||||
dataset_id = payload.get("data", {}).get("dataset_id")
|
||||
assert dataset_id
|
||||
return dataset_id, line
|
||||
|
||||
|
||||
async def test_download_dataset_requires_admin(test_client, standard_user):
|
||||
response = await test_client.get(
|
||||
"/api/evaluation/datasets/dataset_fake/download",
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
async def test_admin_can_download_dataset(test_client, admin_headers, knowledge_database):
|
||||
dataset_id, expected_line = await _upload_test_dataset(test_client, admin_headers, knowledge_database["kb_id"])
|
||||
|
||||
response = await test_client.get(
|
||||
f"/api/evaluation/datasets/{dataset_id}/download",
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert "application/x-ndjson" in response.headers.get("content-type", "")
|
||||
assert "attachment" in response.headers.get("content-disposition", "").lower()
|
||||
|
||||
content = response.content.decode("utf-8")
|
||||
assert expected_line.strip() in content
|
||||
|
||||
|
||||
async def test_download_dataset_not_found(test_client, admin_headers):
|
||||
response = await test_client.get(
|
||||
f"/api/evaluation/datasets/dataset_not_found_{uuid.uuid4().hex[:8]}/download",
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 404, response.text
|
||||
@@ -0,0 +1,30 @@
|
||||
"""
|
||||
Integration tests for graph router list endpoint.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_admin_can_list_graphs(test_client, admin_headers):
|
||||
"""Test that listing graphs returns 200 and a list of graphs."""
|
||||
response = await test_client.get("/api/graph/list", headers=admin_headers)
|
||||
assert response.status_code == 200, f"Failed to list graphs: {response.text}"
|
||||
data = response.json()
|
||||
|
||||
# Check if response is wrapped
|
||||
if isinstance(data, dict) and "data" in data:
|
||||
graphs = data["data"]
|
||||
else:
|
||||
graphs = data
|
||||
|
||||
assert isinstance(graphs, list)
|
||||
# Check structure of returned items if list is not empty
|
||||
if graphs:
|
||||
item = graphs[0]
|
||||
assert "id" in item
|
||||
assert "name" in item
|
||||
assert "type" in item
|
||||
@@ -0,0 +1,781 @@
|
||||
"""
|
||||
Integration tests for knowledge router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from yuxi.knowledge.chunking.ragflow_like.presets import CHUNK_PRESET_IDS
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
def _assert_forbidden_response(response):
|
||||
"""验证 403 禁止访问响应的格式"""
|
||||
assert response.status_code == 403
|
||||
payload = response.json()
|
||||
assert "detail" in payload
|
||||
assert isinstance(payload["detail"], str)
|
||||
|
||||
|
||||
async def _create_test_department(test_client, admin_headers, prefix="pytest_dept"):
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
admin_uid = f"deptadmin_{suffix}"
|
||||
response = await test_client.post(
|
||||
"/api/departments",
|
||||
json={
|
||||
"name": f"{prefix}_{suffix}",
|
||||
"description": "pytest department",
|
||||
"admin_uid": admin_uid,
|
||||
"admin_password": f"Pw!{suffix}",
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 201, response.text
|
||||
payload = response.json()
|
||||
payload["admin_uid"] = admin_uid
|
||||
return payload
|
||||
|
||||
|
||||
async def _create_test_user(test_client, admin_headers, department_id):
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
password = f"Pw!{suffix}"
|
||||
response = await test_client.post(
|
||||
"/api/auth/users",
|
||||
json={
|
||||
"username": f"pytest_user_{suffix}",
|
||||
"password": password,
|
||||
"role": "user",
|
||||
"department_id": department_id,
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
user = response.json()
|
||||
|
||||
login_response = await test_client.post(
|
||||
"/api/auth/token",
|
||||
data={"username": user["uid"], "password": password},
|
||||
)
|
||||
assert login_response.status_code == 200, login_response.text
|
||||
return {"user": user, "headers": {"Authorization": f"Bearer {login_response.json()['access_token']}"}}
|
||||
|
||||
|
||||
async def _delete_user_by_id(test_client, admin_headers, user_id):
|
||||
response = await test_client.delete(f"/api/auth/users/{user_id}", headers=admin_headers)
|
||||
assert response.status_code in (200, 404), response.text
|
||||
|
||||
|
||||
async def _find_user_id_by_uid(test_client, admin_headers, uid):
|
||||
response = await test_client.get("/api/auth/users", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
for user in response.json():
|
||||
if user["uid"] == uid:
|
||||
return user["id"]
|
||||
return None
|
||||
|
||||
|
||||
async def _delete_department_with_admin(test_client, admin_headers, department):
|
||||
admin_user_id = await _find_user_id_by_uid(test_client, admin_headers, department["admin_uid"])
|
||||
if admin_user_id:
|
||||
await _delete_user_by_id(test_client, admin_headers, admin_user_id)
|
||||
response = await test_client.delete(f"/api/departments/{department['id']}", headers=admin_headers)
|
||||
assert response.status_code in (200, 404), response.text
|
||||
|
||||
|
||||
async def _create_test_database(test_client, admin_headers, share_config=None):
|
||||
response = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": f"pytest_acl_{uuid.uuid4().hex[:8]}",
|
||||
"description": "Knowledge permission test",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {},
|
||||
"share_config": share_config,
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()
|
||||
|
||||
|
||||
async def _accessible_kb_ids(test_client, headers):
|
||||
response = await test_client.get("/api/knowledge/databases/accessible", headers=headers)
|
||||
assert response.status_code == 200, response.text
|
||||
return {item["kb_id"] for item in response.json().get("databases", [])}
|
||||
|
||||
|
||||
async def test_admin_can_manage_knowledge_databases(test_client, admin_headers, knowledge_database):
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
list_response = await test_client.get("/api/knowledge/databases", headers=admin_headers)
|
||||
assert list_response.status_code == 200, list_response.text
|
||||
databases = list_response.json().get("databases", [])
|
||||
assert any(entry["kb_id"] == kb_id for entry in databases)
|
||||
|
||||
get_response = await test_client.get(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
assert get_response.status_code == 200, get_response.text
|
||||
assert get_response.json()["kb_id"] == kb_id
|
||||
|
||||
update_response = await test_client.put(
|
||||
f"/api/knowledge/databases/{kb_id}",
|
||||
json={"name": knowledge_database["name"], "description": "Updated by pytest"},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert update_response.status_code == 200, update_response.text
|
||||
assert update_response.json()["database"]["description"] == "Updated by pytest"
|
||||
|
||||
|
||||
async def test_document_exists_returns_false_for_missing_relative_path(test_client, admin_headers, knowledge_database):
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
filename = f"google_drive/shared_drives/engineering/serving-runtime/dsid_{uuid.uuid4().hex}__missing-playbook.txt"
|
||||
|
||||
response = await test_client.get(
|
||||
f"/api/knowledge/databases/{kb_id}/documents/exists",
|
||||
params={"filename": filename},
|
||||
headers=admin_headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"kb_id": kb_id, "filename": filename, "exists": False}
|
||||
|
||||
|
||||
async def test_create_database_with_chunk_preset(test_client, admin_headers):
|
||||
db_name = f"pytest_chunk_preset_{uuid.uuid4().hex[:6]}"
|
||||
payload = {
|
||||
"database_name": db_name,
|
||||
"description": "Chunk preset create test",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {"chunk_preset_id": "book"},
|
||||
}
|
||||
|
||||
create_response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
kb_id = create_response.json()["kb_id"]
|
||||
|
||||
info_response = await test_client.get(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
assert info_response.status_code == 200, info_response.text
|
||||
assert info_response.json()["additional_params"]["chunk_preset_id"] == "book"
|
||||
|
||||
delete_response = await test_client.delete(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
|
||||
|
||||
async def test_get_chunk_presets_returns_configured_options(test_client, admin_headers):
|
||||
response = await test_client.get("/api/knowledge/chunk-presets", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
options = payload["chunk_presets"]
|
||||
assert payload["message"] == "success"
|
||||
assert {option["value"] for option in options} == CHUNK_PRESET_IDS
|
||||
assert all(set(option) == {"value", "label", "description"} for option in options)
|
||||
assert all(option["label"] and option["description"] for option in options)
|
||||
|
||||
|
||||
async def test_update_database_additional_params_merge_keeps_chunk_preset(
|
||||
test_client, admin_headers, knowledge_database
|
||||
):
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
first_update = await test_client.put(
|
||||
f"/api/knowledge/databases/{kb_id}",
|
||||
json={
|
||||
"name": knowledge_database["name"],
|
||||
"description": "update with chunk preset",
|
||||
"additional_params": {"chunk_preset_id": "qa"},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert first_update.status_code == 200, first_update.text
|
||||
|
||||
second_update = await test_client.put(
|
||||
f"/api/knowledge/databases/{kb_id}",
|
||||
json={
|
||||
"name": knowledge_database["name"],
|
||||
"description": "update without additional params",
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert second_update.status_code == 200, second_update.text
|
||||
|
||||
info_response = await test_client.get(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
assert info_response.status_code == 200, info_response.text
|
||||
assert info_response.json()["additional_params"]["chunk_preset_id"] == "qa"
|
||||
|
||||
|
||||
async def test_knowledge_routes_enforce_permissions(test_client, standard_user, knowledge_database):
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
forbidden_create = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": "unauthorized_db",
|
||||
"description": "Should not succeed",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
_assert_forbidden_response(forbidden_create)
|
||||
|
||||
forbidden_list = await test_client.get("/api/knowledge/databases", headers=standard_user["headers"])
|
||||
_assert_forbidden_response(forbidden_list)
|
||||
|
||||
forbidden_chunk_presets = await test_client.get("/api/knowledge/chunk-presets", headers=standard_user["headers"])
|
||||
_assert_forbidden_response(forbidden_chunk_presets)
|
||||
|
||||
forbidden_get = await test_client.get(f"/api/knowledge/databases/{kb_id}", headers=standard_user["headers"])
|
||||
_assert_forbidden_response(forbidden_get)
|
||||
|
||||
forbidden_exists = await test_client.get(
|
||||
f"/api/knowledge/databases/{kb_id}/documents/exists",
|
||||
params={"filename": "demo.txt"},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
_assert_forbidden_response(forbidden_exists)
|
||||
|
||||
|
||||
async def test_admin_can_create_vector_db_with_reranker(test_client, admin_headers):
|
||||
"""测试创建向量库并配置 reranker 参数(通过 query_params.options)
|
||||
|
||||
注意:数据库清理由 conftest.py 中的 session fixture 自动处理。
|
||||
"""
|
||||
db_name = f"pytest_rerank_{uuid.uuid4().hex[:6]}"
|
||||
payload = {
|
||||
"database_name": db_name,
|
||||
"description": "Vector DB with reranker",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {},
|
||||
}
|
||||
|
||||
create_response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
|
||||
db_payload = create_response.json()
|
||||
kb_id = db_payload["kb_id"]
|
||||
|
||||
# 获取查询参数配置
|
||||
params_response = await test_client.get(f"/api/knowledge/databases/{kb_id}/query-params", headers=admin_headers)
|
||||
assert params_response.status_code == 200, params_response.text
|
||||
|
||||
params_payload = params_response.json()
|
||||
options = params_payload.get("params", {}).get("options", [])
|
||||
option_keys = {option.get("key") for option in options}
|
||||
|
||||
# 验证新的参数名称
|
||||
assert "final_top_k" in option_keys
|
||||
assert "use_reranker" in option_keys
|
||||
assert "recall_top_k" in option_keys
|
||||
assert "reranker_model" in option_keys
|
||||
|
||||
# 验证参数配置
|
||||
final_top_k_option = next((opt for opt in options if opt.get("key") == "final_top_k"), None)
|
||||
assert final_top_k_option is not None
|
||||
assert final_top_k_option.get("default") == 10
|
||||
|
||||
use_reranker_option = next((opt for opt in options if opt.get("key") == "use_reranker"), None)
|
||||
assert use_reranker_option is not None
|
||||
assert use_reranker_option.get("default") is False
|
||||
|
||||
# 保存查询参数(模拟前端配置)
|
||||
update_params = {
|
||||
"final_top_k": 5,
|
||||
"use_reranker": True,
|
||||
"recall_top_k": 20,
|
||||
}
|
||||
update_response = await test_client.put(
|
||||
f"/api/knowledge/databases/{kb_id}/query-params", json=update_params, headers=admin_headers
|
||||
)
|
||||
assert update_response.status_code == 200, update_response.text
|
||||
|
||||
# 再次获取参数,验证保存成功
|
||||
params_response2 = await test_client.get(f"/api/knowledge/databases/{kb_id}/query-params", headers=admin_headers)
|
||||
assert params_response2.status_code == 200, params_response2.text
|
||||
|
||||
params_payload2 = params_response2.json()
|
||||
options2 = params_payload2.get("params", {}).get("options", [])
|
||||
|
||||
# 验证保存的值
|
||||
final_top_k_option2 = next((opt for opt in options2 if opt.get("key") == "final_top_k"), None)
|
||||
assert final_top_k_option2 is not None
|
||||
assert final_top_k_option2.get("default") == 5 # 保存的值
|
||||
|
||||
use_reranker_option2 = next((opt for opt in options2 if opt.get("key") == "use_reranker"), None)
|
||||
assert use_reranker_option2 is not None
|
||||
assert use_reranker_option2.get("default") is True # 保存的值
|
||||
|
||||
|
||||
async def test_create_dify_database_success(test_client, admin_headers):
|
||||
db_name = f"pytest_dify_{uuid.uuid4().hex[:6]}"
|
||||
payload = {
|
||||
"database_name": db_name,
|
||||
"description": "Dify KB create test",
|
||||
"kb_type": "dify",
|
||||
"additional_params": {
|
||||
"dify_api_url": "https://api.dify.ai/v1",
|
||||
"dify_token": "test-token",
|
||||
"dify_dataset_id": "dataset-123",
|
||||
},
|
||||
}
|
||||
|
||||
create_response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
created_payload = create_response.json()
|
||||
kb_id = created_payload["kb_id"]
|
||||
assert created_payload["embedding_model_spec"] is None
|
||||
assert "chunk_preset_id" not in created_payload["metadata"]
|
||||
|
||||
info_response = await test_client.get(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
assert info_response.status_code == 200, info_response.text
|
||||
additional_params = info_response.json()["additional_params"]
|
||||
assert additional_params["dify_api_url"] == "https://api.dify.ai/v1"
|
||||
assert additional_params["dify_token"] == "test-token"
|
||||
assert additional_params["dify_dataset_id"] == "dataset-123"
|
||||
|
||||
|
||||
async def test_create_dify_database_missing_params_failed(test_client, admin_headers):
|
||||
payload = {
|
||||
"database_name": f"pytest_dify_missing_{uuid.uuid4().hex[:6]}",
|
||||
"description": "Dify KB missing params",
|
||||
"kb_type": "dify",
|
||||
"additional_params": {
|
||||
"dify_api_url": "https://api.dify.ai/v1",
|
||||
"dify_token": "",
|
||||
"dify_dataset_id": "",
|
||||
},
|
||||
}
|
||||
|
||||
response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "Dify 参数缺失" in response.json()["detail"]
|
||||
|
||||
|
||||
async def test_create_dify_database_invalid_api_url_failed(test_client, admin_headers):
|
||||
payload = {
|
||||
"database_name": f"pytest_dify_bad_url_{uuid.uuid4().hex[:6]}",
|
||||
"description": "Dify KB invalid api url",
|
||||
"kb_type": "dify",
|
||||
"additional_params": {
|
||||
"dify_api_url": "https://api.dify.ai",
|
||||
"dify_token": "test-token",
|
||||
"dify_dataset_id": "dataset-123",
|
||||
},
|
||||
}
|
||||
|
||||
response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert response.status_code == 400, response.text
|
||||
assert "/v1" in response.json()["detail"]
|
||||
|
||||
|
||||
async def test_dify_query_params_and_documents_readonly(test_client, admin_headers):
|
||||
payload = {
|
||||
"database_name": f"pytest_dify_ro_{uuid.uuid4().hex[:6]}",
|
||||
"description": "Dify readonly routes",
|
||||
"kb_type": "dify",
|
||||
"additional_params": {
|
||||
"dify_api_url": "https://api.dify.ai/v1",
|
||||
"dify_token": "test-token",
|
||||
"dify_dataset_id": "dataset-123",
|
||||
},
|
||||
}
|
||||
|
||||
create_response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
kb_id = create_response.json()["kb_id"]
|
||||
|
||||
params_response = await test_client.get(f"/api/knowledge/databases/{kb_id}/query-params", headers=admin_headers)
|
||||
assert params_response.status_code == 200, params_response.text
|
||||
options = params_response.json().get("params", {}).get("options", [])
|
||||
option_keys = {item.get("key") for item in options}
|
||||
assert option_keys == {"search_mode", "final_top_k", "score_threshold_enabled", "similarity_threshold"}
|
||||
|
||||
add_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/documents",
|
||||
json={"items": ["/tmp/demo.txt"], "params": {"content_type": "file"}},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert add_response.status_code == 400, add_response.text
|
||||
assert "只支持检索" in add_response.json()["detail"]
|
||||
|
||||
parse_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/documents/parse",
|
||||
json=["file_id_1"],
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert parse_response.status_code == 400, parse_response.text
|
||||
assert "只支持检索" in parse_response.json()["detail"]
|
||||
|
||||
index_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/documents/index",
|
||||
json={"file_ids": ["file_id_1"], "params": {}},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert index_response.status_code == 400, index_response.text
|
||||
assert "只支持检索" in index_response.json()["detail"]
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# === Mindmap Tests ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_get_databases_overview(test_client, admin_headers, knowledge_database):
|
||||
"""测试获取所有知识库概览"""
|
||||
response = await test_client.get("/api/knowledge/mindmap/databases", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert "databases" in payload
|
||||
assert "total" in payload
|
||||
|
||||
# 验证知识库在列表中
|
||||
kb_ids = [db["kb_id"] for db in payload["databases"]]
|
||||
assert knowledge_database["kb_id"] in kb_ids
|
||||
|
||||
|
||||
async def test_get_database_files(test_client, admin_headers, knowledge_database):
|
||||
"""测试获取知识库文件列表"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
response = await test_client.get(f"/api/knowledge/databases/{kb_id}/mindmap/files", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert payload["kb_id"] == kb_id
|
||||
assert "files" in payload
|
||||
assert "total" in payload
|
||||
assert payload["db_name"] == knowledge_database["name"]
|
||||
|
||||
|
||||
async def test_get_database_files_not_found(test_client, admin_headers):
|
||||
"""测试获取不存在的知识库文件列表"""
|
||||
response = await test_client.get("/api/knowledge/databases/nonexistent_kb_id/mindmap/files", headers=admin_headers)
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
async def test_generate_mindmap_empty_files(test_client, admin_headers, knowledge_database):
|
||||
"""测试空文件列表生成思维导图"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/mindmap/generate",
|
||||
json={"file_ids": [], "user_prompt": ""},
|
||||
headers=admin_headers,
|
||||
)
|
||||
# 空文件应该返回400错误
|
||||
assert response.status_code == 400
|
||||
assert "中没有文件" in response.json()["detail"]
|
||||
|
||||
|
||||
async def test_get_database_mindmap_not_exists(test_client, admin_headers, knowledge_database):
|
||||
"""测试获取不存在的思维导图"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
response = await test_client.get(f"/api/knowledge/databases/{kb_id}/mindmap", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["kb_id"] == kb_id
|
||||
assert payload["mindmap"] is None # 尚未生成思维导图
|
||||
|
||||
|
||||
async def test_generate_and_get_mindmap(test_client, admin_headers, knowledge_database):
|
||||
"""测试生成并获取思维导图
|
||||
|
||||
注意:此测试需要知识库中有文件才能完整测试核心功能。
|
||||
由于没有前置的文件上传 fixture,测试会先验证空文件场景(预期400),
|
||||
然后使用 xfail 标记等待后续完善。
|
||||
"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
# 空文件场景 - 预期返回400错误
|
||||
generate_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/mindmap/generate",
|
||||
json={"file_ids": [], "user_prompt": ""},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert generate_response.status_code == 400
|
||||
assert "中没有文件" in generate_response.json()["detail"]
|
||||
|
||||
# 标记此测试需要文件上传支持才能完整执行
|
||||
pytest.skip("需要先上传文件才能完整测试思维导图生成功能")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# === Knowledge Router Additional Tests ===
|
||||
# =============================================================================
|
||||
|
||||
|
||||
async def test_get_accessible_databases(test_client, admin_headers, knowledge_database):
|
||||
"""测试获取可访问的知识库列表"""
|
||||
response = await test_client.get("/api/knowledge/databases/accessible", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert "databases" in payload
|
||||
|
||||
# 验证知识库在列表中
|
||||
kb_ids = [db["kb_id"] for db in payload["databases"]]
|
||||
assert knowledge_database["kb_id"] in kb_ids
|
||||
|
||||
|
||||
async def test_create_database_defaults_to_global_share_config(test_client, admin_headers):
|
||||
database = await _create_test_database(test_client, admin_headers)
|
||||
kb_id = database["kb_id"]
|
||||
try:
|
||||
assert database["share_config"] == {"access_level": "global", "department_ids": [], "user_uids": []}
|
||||
finally:
|
||||
await test_client.delete(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
|
||||
|
||||
async def test_department_share_config_filters_accessible_databases(test_client, admin_headers):
|
||||
department_a = await _create_test_department(test_client, admin_headers, "pytest_dept_a")
|
||||
department_b = await _create_test_department(test_client, admin_headers, "pytest_dept_b")
|
||||
user_a = user_b = None
|
||||
database = None
|
||||
|
||||
try:
|
||||
user_a = await _create_test_user(test_client, admin_headers, department_a["id"])
|
||||
user_b = await _create_test_user(test_client, admin_headers, department_b["id"])
|
||||
database = await _create_test_database(
|
||||
test_client,
|
||||
admin_headers,
|
||||
{"access_level": "department", "department_ids": [department_a["id"]], "user_uids": []},
|
||||
)
|
||||
|
||||
saved_config = database["share_config"]
|
||||
assert saved_config["access_level"] == "department"
|
||||
assert department_a["id"] in saved_config["department_ids"]
|
||||
|
||||
assert database["kb_id"] in await _accessible_kb_ids(test_client, user_a["headers"])
|
||||
assert database["kb_id"] not in await _accessible_kb_ids(test_client, user_b["headers"])
|
||||
finally:
|
||||
if database:
|
||||
await test_client.delete(f"/api/knowledge/databases/{database['kb_id']}", headers=admin_headers)
|
||||
if user_a:
|
||||
await _delete_user_by_id(test_client, admin_headers, user_a["user"]["id"])
|
||||
if user_b:
|
||||
await _delete_user_by_id(test_client, admin_headers, user_b["user"]["id"])
|
||||
await _delete_department_with_admin(test_client, admin_headers, department_a)
|
||||
await _delete_department_with_admin(test_client, admin_headers, department_b)
|
||||
|
||||
|
||||
async def test_user_share_config_filters_accessible_databases(test_client, admin_headers):
|
||||
department_a = await _create_test_department(test_client, admin_headers, "pytest_dept_a")
|
||||
department_b = await _create_test_department(test_client, admin_headers, "pytest_dept_b")
|
||||
user_a = user_b = None
|
||||
database = None
|
||||
|
||||
try:
|
||||
user_a = await _create_test_user(test_client, admin_headers, department_a["id"])
|
||||
user_b = await _create_test_user(test_client, admin_headers, department_b["id"])
|
||||
database = await _create_test_database(
|
||||
test_client,
|
||||
admin_headers,
|
||||
{"access_level": "user", "department_ids": [], "user_uids": [user_a["user"]["uid"]]},
|
||||
)
|
||||
|
||||
saved_config = database["share_config"]
|
||||
assert saved_config["access_level"] == "user"
|
||||
assert user_a["user"]["uid"] in saved_config["user_uids"]
|
||||
|
||||
assert database["kb_id"] in await _accessible_kb_ids(test_client, user_a["headers"])
|
||||
assert database["kb_id"] not in await _accessible_kb_ids(test_client, user_b["headers"])
|
||||
finally:
|
||||
if database:
|
||||
await test_client.delete(f"/api/knowledge/databases/{database['kb_id']}", headers=admin_headers)
|
||||
if user_a:
|
||||
await _delete_user_by_id(test_client, admin_headers, user_a["user"]["id"])
|
||||
if user_b:
|
||||
await _delete_user_by_id(test_client, admin_headers, user_b["user"]["id"])
|
||||
await _delete_department_with_admin(test_client, admin_headers, department_a)
|
||||
await _delete_department_with_admin(test_client, admin_headers, department_b)
|
||||
|
||||
|
||||
async def test_user_access_options_include_all_departments_for_admin(test_client, admin_headers):
|
||||
department = await _create_test_department(test_client, admin_headers, "pytest_access_options")
|
||||
user = None
|
||||
|
||||
try:
|
||||
user = await _create_test_user(test_client, admin_headers, department["id"])
|
||||
response = await test_client.get("/api/auth/users/access-options", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
uids = {item["uid"] for item in response.json()}
|
||||
assert user["user"]["uid"] in uids
|
||||
assert department["admin_uid"] in uids
|
||||
finally:
|
||||
if user:
|
||||
await _delete_user_by_id(test_client, admin_headers, user["user"]["id"])
|
||||
await _delete_department_with_admin(test_client, admin_headers, department)
|
||||
|
||||
|
||||
async def test_get_knowledge_base_types(test_client, admin_headers):
|
||||
"""测试获取支持的知识库类型"""
|
||||
response = await test_client.get("/api/knowledge/types", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert "kb_types" in payload
|
||||
assert "default_config" not in payload["kb_types"]["dify"]
|
||||
assert payload["kb_types"]["dify"]["name"] == "Dify"
|
||||
assert payload["kb_types"]["dify"]["description"] == "连接 Dify Dataset 的只读检索知识库"
|
||||
assert payload["kb_types"]["dify"]["requires_embedding_model"] is False
|
||||
assert payload["kb_types"]["dify"]["supports_documents"] is False
|
||||
assert [option["key"] for option in payload["kb_types"]["dify"]["create_params"]["options"]] == [
|
||||
"dify_api_url",
|
||||
"dify_token",
|
||||
"dify_dataset_id",
|
||||
]
|
||||
assert "default_config" not in payload["kb_types"]["notion"]
|
||||
assert payload["kb_types"]["notion"]["name"] == "Notion"
|
||||
assert (
|
||||
payload["kb_types"]["notion"]["description"]
|
||||
== "连接 Notion Data Source 的只读知识库,支持检索、打开页面和页内查找"
|
||||
)
|
||||
assert payload["kb_types"]["notion"]["requires_embedding_model"] is False
|
||||
assert payload["kb_types"]["notion"]["supports_documents"] is False
|
||||
assert [option["key"] for option in payload["kb_types"]["notion"]["create_params"]["options"]] == [
|
||||
"notion_token",
|
||||
"notion_data_source_id",
|
||||
"notion_version",
|
||||
]
|
||||
|
||||
|
||||
async def test_get_knowledge_base_statistics(test_client, admin_headers):
|
||||
"""测试获取知识库统计信息"""
|
||||
response = await test_client.get("/api/knowledge/stats", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert "stats" in payload
|
||||
|
||||
|
||||
async def test_get_supported_file_types(test_client, admin_headers):
|
||||
"""测试获取支持的文件类型"""
|
||||
response = await test_client.get("/api/knowledge/files/supported-types", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert "file_types" in payload
|
||||
assert isinstance(payload["file_types"], list)
|
||||
|
||||
|
||||
async def test_markdown_endpoint_parses_uploaded_text_file(test_client, admin_headers):
|
||||
"""测试 /files/markdown 能解析上传文件并返回 markdown。"""
|
||||
data_dir = Path(__file__).resolve().parents[2] / "data"
|
||||
test_file = data_dir / "A_Dream_of_Red_Mansions_10hui.txt"
|
||||
|
||||
assert test_file.exists(), f"测试文件不存在: {test_file}"
|
||||
|
||||
with test_file.open("rb") as f:
|
||||
response = await test_client.post(
|
||||
"/api/knowledge/files/markdown",
|
||||
headers=admin_headers,
|
||||
files={"file": (test_file.name, f, "text/plain")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["message"] == "success"
|
||||
assert isinstance(payload.get("markdown_content"), str)
|
||||
assert payload["markdown_content"].strip()
|
||||
|
||||
|
||||
async def test_duplicate_database_name(test_client, admin_headers, knowledge_database):
|
||||
"""测试重复创建同名知识库"""
|
||||
db_name = knowledge_database["name"]
|
||||
response = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": db_name,
|
||||
"description": "Duplicate name test",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 409
|
||||
assert "已存在" in response.json()["detail"]
|
||||
|
||||
|
||||
async def test_create_lightrag_knowledge_base_is_unsupported(test_client, admin_headers):
|
||||
db_name = f"pytest_lightrag_{uuid.uuid4().hex[:6]}"
|
||||
response = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": db_name,
|
||||
"description": "Unsupported LightRAG knowledge base",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "lightrag",
|
||||
"additional_params": {},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert "Unsupported knowledge base type: lightrag" in response.json()["detail"]
|
||||
|
||||
|
||||
async def test_create_milvus_knowledge_base(test_client, admin_headers):
|
||||
"""测试创建 Milvus 知识库
|
||||
|
||||
注意:数据库清理由 conftest.py 中的 session fixture 自动处理。
|
||||
"""
|
||||
db_name = f"pytest_milvus_{uuid.uuid4().hex[:6]}"
|
||||
payload = {
|
||||
"database_name": db_name,
|
||||
"description": "Pytest Milvus knowledge base",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {},
|
||||
}
|
||||
|
||||
create_response = await test_client.post("/api/knowledge/databases", json=payload, headers=admin_headers)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
|
||||
db_payload = create_response.json()
|
||||
assert db_payload["kb_type"] == "milvus"
|
||||
|
||||
|
||||
async def test_sample_questions_endpoints(test_client, admin_headers, knowledge_database):
|
||||
"""测试示例问题接口(空文件时预期返回400)"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
# 获取示例问题(空知识库应该返回空列表)
|
||||
get_response = await test_client.get(f"/api/knowledge/databases/{kb_id}/sample-questions", headers=admin_headers)
|
||||
assert get_response.status_code == 200, get_response.text
|
||||
get_payload = get_response.json()
|
||||
assert get_payload["kb_id"] == kb_id
|
||||
assert "questions" in get_payload
|
||||
assert get_payload["count"] == 0 # 空知识库没有问题
|
||||
|
||||
# 生成示例问题(空知识库应该返回400)
|
||||
generate_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/sample-questions",
|
||||
json={"count": 5},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert generate_response.status_code == 400
|
||||
assert "中没有文件" in generate_response.json()["detail"]
|
||||
|
||||
|
||||
async def test_mindmap_permissions(test_client, standard_user, knowledge_database):
|
||||
"""测试思维导图接口的权限控制"""
|
||||
kb_id = knowledge_database["kb_id"]
|
||||
|
||||
# 普通用户应该无法访问
|
||||
forbidden_list = await test_client.get("/api/knowledge/mindmap/databases", headers=standard_user["headers"])
|
||||
_assert_forbidden_response(forbidden_list)
|
||||
|
||||
forbidden_files = await test_client.get(
|
||||
f"/api/knowledge/databases/{kb_id}/mindmap/files", headers=standard_user["headers"]
|
||||
)
|
||||
_assert_forbidden_response(forbidden_files)
|
||||
|
||||
forbidden_generate = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/mindmap/generate",
|
||||
json={"file_ids": []},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
_assert_forbidden_response(forbidden_generate)
|
||||
@@ -0,0 +1,22 @@
|
||||
"""
|
||||
Integration tests for settings router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_reranker_list_requires_admin(test_client, standard_user):
|
||||
public_response = await test_client.get("/api/settings/rerankers")
|
||||
assert public_response.status_code == 404
|
||||
|
||||
forbidden_response = await test_client.get("/api/settings/rerankers", headers=standard_user["headers"])
|
||||
assert forbidden_response.status_code == 404
|
||||
|
||||
|
||||
async def test_admin_can_list_rerankers(test_client, admin_headers):
|
||||
response = await test_client.get("/api/settings/rerankers", headers=admin_headers)
|
||||
assert response.status_code == 404, response.text
|
||||
@@ -0,0 +1,59 @@
|
||||
"""
|
||||
Integration tests for system router endpoints.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def test_health_endpoint_is_public(test_client):
|
||||
response = await test_client.get("/api/system/health")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "ok"
|
||||
|
||||
|
||||
async def test_info_endpoint_is_public(test_client):
|
||||
response = await test_client.get("/api/system/info")
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert "data" in payload
|
||||
|
||||
|
||||
async def test_config_get_requires_login_and_update_requires_admin(test_client, standard_user):
|
||||
assert (await test_client.get("/api/system/config")).status_code == 401
|
||||
user_config_response = await test_client.get("/api/system/config", headers=standard_user["headers"])
|
||||
assert user_config_response.status_code == 200, user_config_response.text
|
||||
|
||||
update_response = await test_client.post(
|
||||
"/api/system/config/update",
|
||||
json={"default_ocr_engine": "rapid_ocr"},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
assert update_response.status_code == 403
|
||||
|
||||
|
||||
async def test_admin_can_fetch_config_and_reload_info(test_client, admin_headers):
|
||||
config_response = await test_client.get("/api/system/config", headers=admin_headers)
|
||||
assert config_response.status_code == 200, config_response.text
|
||||
assert isinstance(config_response.json(), dict)
|
||||
|
||||
reload_response = await test_client.post("/api/system/info/reload", headers=admin_headers)
|
||||
assert reload_response.status_code == 200, reload_response.text
|
||||
reload_payload = reload_response.json()
|
||||
assert reload_payload["success"] is True
|
||||
assert "data" in reload_payload
|
||||
|
||||
|
||||
async def test_admin_can_fetch_tools_with_config_guide_field(test_client, admin_headers):
|
||||
response = await test_client.get("/api/system/tools", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert isinstance(payload["data"], list)
|
||||
assert payload["data"]
|
||||
assert "config_guide" in payload["data"][0]
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
Integration tests for the task management router.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
_LITE_MODE = os.getenv("LITE_MODE", "").lower() in {"true", "1"}
|
||||
|
||||
|
||||
async def test_task_routes_require_admin(test_client, standard_user):
|
||||
"""Non-admin users should be blocked from accessing task APIs."""
|
||||
headers = standard_user["headers"]
|
||||
|
||||
list_response = await test_client.get("/api/tasks", headers=headers)
|
||||
assert list_response.status_code == 403
|
||||
|
||||
detail_response = await test_client.get("/api/tasks/some-task", headers=headers)
|
||||
assert detail_response.status_code == 403
|
||||
|
||||
cancel_response = await test_client.post("/api/tasks/some-task/cancel", headers=headers)
|
||||
assert cancel_response.status_code == 403
|
||||
|
||||
|
||||
async def test_admin_can_list_tasks(test_client, admin_headers):
|
||||
"""Admin should receive a well-formed task list payload."""
|
||||
response = await test_client.get("/api/tasks", headers=admin_headers)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
assert "tasks" in payload
|
||||
assert isinstance(payload["tasks"], list)
|
||||
assert "summary" in payload
|
||||
assert isinstance(payload["summary"], dict)
|
||||
|
||||
|
||||
async def test_cancel_unknown_task_returns_client_error(test_client, admin_headers):
|
||||
"""Cancelling a non-existent task should surface a 400 response."""
|
||||
response = await test_client.post("/api/tasks/not-real/cancel", headers=admin_headers)
|
||||
assert response.status_code == 400, response.text
|
||||
|
||||
|
||||
async def test_enqueue_document_creates_task(
|
||||
test_client,
|
||||
admin_headers,
|
||||
):
|
||||
"""Trigger knowledge ingestion to ensure a task record is materialised."""
|
||||
if _LITE_MODE:
|
||||
enqueue_response = await test_client.post(
|
||||
"/api/knowledge/databases/lite-mode-disabled/documents",
|
||||
json={"items": [], "params": {"content_type": "file"}},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert enqueue_response.status_code == 404
|
||||
return
|
||||
|
||||
create_response = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": f"pytest_task_router_{uuid.uuid4().hex[:8]}",
|
||||
"description": "Task router integration test",
|
||||
"embedding_model_spec": "siliconflow-cn:Pro/BAAI/bge-m3",
|
||||
"kb_type": "milvus",
|
||||
"additional_params": {},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert create_response.status_code == 200, create_response.text
|
||||
kb_id = create_response.json()["kb_id"]
|
||||
|
||||
try:
|
||||
enqueue_response = await test_client.post(
|
||||
f"/api/knowledge/databases/{kb_id}/documents",
|
||||
json={
|
||||
"items": [],
|
||||
"params": {"content_type": "file"},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert enqueue_response.status_code == 200, enqueue_response.text
|
||||
|
||||
enqueue_payload = enqueue_response.json()
|
||||
assert enqueue_payload.get("status") == "queued"
|
||||
task_id = enqueue_payload.get("task_id")
|
||||
assert task_id, "Knowledge ingestion did not return a task_id"
|
||||
|
||||
# The task should be queryable immediately after enqueueing.
|
||||
detail_response = await test_client.get(f"/api/tasks/{task_id}", headers=admin_headers)
|
||||
assert detail_response.status_code == 200, detail_response.text
|
||||
detail_payload = detail_response.json().get("task", {})
|
||||
assert detail_payload.get("id") == task_id
|
||||
assert detail_payload.get("status") in {"queued", "pending", "running", "failed", "success", "cancelled"}
|
||||
|
||||
# Ensure the task surfaces in the list endpoint within a short window.
|
||||
for _ in range(10):
|
||||
list_response = await test_client.get("/api/tasks", headers=admin_headers)
|
||||
assert list_response.status_code == 200, list_response.text
|
||||
all_tasks = list_response.json().get("tasks", [])
|
||||
if any(entry.get("id") == task_id for entry in all_tasks):
|
||||
break
|
||||
await asyncio.sleep(0.2)
|
||||
else:
|
||||
pytest.fail("Task did not appear in list endpoint within timeout window")
|
||||
|
||||
# Poll for terminal state to validate worker bookkeeping.
|
||||
for _ in range(20):
|
||||
detail_response = await test_client.get(f"/api/tasks/{task_id}", headers=admin_headers)
|
||||
task_status = detail_response.json().get("task", {}).get("status")
|
||||
if task_status in {"success", "failed", "cancelled"}:
|
||||
break
|
||||
await asyncio.sleep(0.5)
|
||||
else:
|
||||
pytest.fail("Task did not reach a terminal status within timeout window")
|
||||
finally:
|
||||
await test_client.delete(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
@@ -0,0 +1,123 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def _create_dify_database(test_client, admin_headers) -> str:
|
||||
response = await test_client.post(
|
||||
"/api/knowledge/databases",
|
||||
json={
|
||||
"database_name": f"pytest_graph_dify_{uuid.uuid4().hex[:8]}",
|
||||
"description": "Graph router Dify negative test",
|
||||
"kb_type": "dify",
|
||||
"additional_params": {
|
||||
"dify_api_url": "https://api.dify.ai/v1",
|
||||
"dify_token": "test-token",
|
||||
"dify_dataset_id": f"dataset-{uuid.uuid4().hex[:8]}",
|
||||
},
|
||||
},
|
||||
headers=admin_headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return response.json()["kb_id"]
|
||||
|
||||
|
||||
async def _delete_database(test_client, admin_headers, kb_id: str) -> None:
|
||||
await test_client.delete(f"/api/knowledge/databases/{kb_id}", headers=admin_headers)
|
||||
|
||||
|
||||
async def test_graph_routes_require_auth(test_client):
|
||||
response = await test_client.get("/api/graph/list")
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_standard_user_cannot_access_graph_endpoints(test_client, standard_user):
|
||||
response = await test_client.get("/api/graph/list", headers=standard_user["headers"])
|
||||
assert response.status_code == 403
|
||||
|
||||
|
||||
async def test_get_graphs_list_only_returns_milvus(test_client, admin_headers, knowledge_database):
|
||||
response = await test_client.get("/api/graph/list", headers=admin_headers)
|
||||
|
||||
assert response.status_code == 200
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert isinstance(payload["data"], list)
|
||||
assert payload["data"]
|
||||
assert all(graph["type"] == "milvus" for graph in payload["data"])
|
||||
assert any(graph["id"] == knowledge_database["kb_id"] for graph in payload["data"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", ["/api/graph/subgraph", "/api/graph/stats", "/api/graph/labels"])
|
||||
async def test_graph_endpoints_reject_non_milvus_types(test_client, admin_headers, path):
|
||||
kb_id = await _create_dify_database(test_client, admin_headers)
|
||||
try:
|
||||
response = await test_client.get(path, params={"kb_id": kb_id}, headers=admin_headers)
|
||||
finally:
|
||||
await _delete_database(test_client, admin_headers, kb_id)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert "only supports Milvus" in response.text
|
||||
|
||||
|
||||
async def test_milvus_subgraph_endpoint(test_client, admin_headers, knowledge_database):
|
||||
response = await test_client.get(
|
||||
"/api/graph/subgraph",
|
||||
params={"kb_id": knowledge_database["kb_id"], "node_label": "*", "max_nodes": 10},
|
||||
headers=admin_headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert "nodes" in payload["data"]
|
||||
assert "edges" in payload["data"]
|
||||
|
||||
|
||||
async def test_milvus_stats_endpoint(test_client, admin_headers, knowledge_database):
|
||||
response = await test_client.get(
|
||||
"/api/graph/stats",
|
||||
params={"kb_id": knowledge_database["kb_id"]},
|
||||
headers=admin_headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert "total_nodes" in payload["data"]
|
||||
assert "total_edges" in payload["data"]
|
||||
assert "entity_types" in payload["data"]
|
||||
|
||||
|
||||
async def test_milvus_labels_endpoint(test_client, admin_headers, knowledge_database):
|
||||
response = await test_client.get(
|
||||
"/api/graph/labels",
|
||||
params={"kb_id": knowledge_database["kb_id"]},
|
||||
headers=admin_headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert "labels" in payload["data"]
|
||||
assert isinstance(payload["data"]["labels"], list)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
[
|
||||
"/api/graph/neo4j/nodes",
|
||||
"/api/graph/neo4j/node",
|
||||
"/api/graph/neo4j/info",
|
||||
"/api/graph/neo4j/index-entities",
|
||||
"/api/graph/neo4j/add-entities",
|
||||
],
|
||||
)
|
||||
async def test_neo4j_upload_routes_are_removed(test_client, admin_headers, path):
|
||||
method = test_client.post if path.endswith(("index-entities", "add-entities")) else test_client.get
|
||||
response = await method(path, headers=admin_headers)
|
||||
assert response.status_code == 404
|
||||
@@ -0,0 +1,71 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
AGENT_ENV_PATH = "/api/user/agent-env"
|
||||
|
||||
|
||||
async def test_agent_env_requires_auth(test_client):
|
||||
response = await test_client.get(AGENT_ENV_PATH)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_agent_env_round_trip_and_replace(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
|
||||
initial_response = await test_client.get(AGENT_ENV_PATH, headers=headers)
|
||||
assert initial_response.status_code == 200, initial_response.text
|
||||
assert initial_response.json()["env"] == {}
|
||||
|
||||
payload = {"env": {"YUXI_TEST_TOKEN": "secret", "YUXI_EMPTY_VALUE": ""}}
|
||||
save_response = await test_client.put(AGENT_ENV_PATH, json=payload, headers=headers)
|
||||
assert save_response.status_code == 200, save_response.text
|
||||
assert save_response.json()["env"] == payload["env"]
|
||||
|
||||
replace_response = await test_client.put(
|
||||
AGENT_ENV_PATH,
|
||||
json={"env": {"YUXI_TEST_TOKEN": "updated"}},
|
||||
headers=headers,
|
||||
)
|
||||
assert replace_response.status_code == 200, replace_response.text
|
||||
assert replace_response.json()["env"] == {"YUXI_TEST_TOKEN": "updated"}
|
||||
|
||||
final_response = await test_client.get(AGENT_ENV_PATH, headers=headers)
|
||||
assert final_response.status_code == 200, final_response.text
|
||||
assert final_response.json()["env"] == {"YUXI_TEST_TOKEN": "updated"}
|
||||
|
||||
|
||||
async def test_agent_env_rejects_invalid_keys(test_client, standard_user):
|
||||
response = await test_client.put(
|
||||
AGENT_ENV_PATH,
|
||||
json={"env": {"INVALID-KEY": "value"}},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
async def test_agent_env_rejects_duplicate_normalized_keys(test_client, standard_user):
|
||||
response = await test_client.put(
|
||||
AGENT_ENV_PATH,
|
||||
json={"env": {"YUXI_TOKEN": "first", " YUXI_TOKEN ": "second"}},
|
||||
headers=standard_user["headers"],
|
||||
)
|
||||
assert response.status_code == 400
|
||||
|
||||
|
||||
async def test_agent_env_is_user_scoped(test_client, standard_user, admin_headers):
|
||||
standard_headers = standard_user["headers"]
|
||||
user_payload = {"env": {"YUXI_USER_ONLY": "user-value"}}
|
||||
|
||||
user_save_response = await test_client.put(AGENT_ENV_PATH, json=user_payload, headers=standard_headers)
|
||||
assert user_save_response.status_code == 200, user_save_response.text
|
||||
|
||||
admin_response = await test_client.get(AGENT_ENV_PATH, headers=admin_headers)
|
||||
assert admin_response.status_code == 200, admin_response.text
|
||||
assert admin_response.json()["env"] != user_payload["env"]
|
||||
|
||||
user_response = await test_client.get(AGENT_ENV_PATH, headers=standard_headers)
|
||||
assert user_response.status_code == 200, user_response.text
|
||||
assert user_response.json()["env"] == user_payload["env"]
|
||||
@@ -0,0 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
USER_CONFIG_PATH = "/api/user/config"
|
||||
|
||||
|
||||
async def test_user_config_requires_auth(test_client):
|
||||
response = await test_client.get(USER_CONFIG_PATH)
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_user_config_round_trip(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
|
||||
initial_response = await test_client.get(USER_CONFIG_PATH, headers=headers)
|
||||
assert initial_response.status_code == 200, initial_response.text
|
||||
initial_payload = initial_response.json()
|
||||
assert initial_payload["enable_memory"] is False
|
||||
|
||||
save_response = await test_client.put(USER_CONFIG_PATH, json={"enable_memory": True}, headers=headers)
|
||||
assert save_response.status_code == 200, save_response.text
|
||||
assert save_response.json()["enable_memory"] is True
|
||||
|
||||
final_response = await test_client.get(USER_CONFIG_PATH, headers=headers)
|
||||
assert final_response.status_code == 200, final_response.text
|
||||
assert final_response.json()["enable_memory"] is True
|
||||
@@ -0,0 +1,690 @@
|
||||
"""Integration tests for viewer-oriented filesystem router endpoints."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
|
||||
from fastapi import HTTPException
|
||||
import pytest
|
||||
from yuxi.agents.backends.sandbox import (
|
||||
ensure_thread_dirs,
|
||||
sandbox_user_data_dir,
|
||||
sandbox_workspace_dir,
|
||||
virtual_path_for_thread_file,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
def _get_provider():
|
||||
from yuxi.agents.backends import get_sandbox_provider
|
||||
|
||||
return get_sandbox_provider()
|
||||
|
||||
|
||||
async def _create_thread_for_user(test_client, headers: dict[str, str]) -> str:
|
||||
agents_resp = await test_client.get("/api/agent", headers=headers)
|
||||
assert agents_resp.status_code == 200, agents_resp.text
|
||||
agents = agents_resp.json().get("agents", [])
|
||||
if not agents:
|
||||
pytest.skip("No agents available for viewer filesystem integration tests.")
|
||||
|
||||
agent_id = agents[0].get("agent_id") or agents[0].get("slug")
|
||||
if not agent_id:
|
||||
pytest.skip("Agent payload missing agent_id field.")
|
||||
|
||||
create_resp = await test_client.post(
|
||||
"/api/chat/thread",
|
||||
json={
|
||||
"agent_id": agent_id,
|
||||
"title": f"viewer-filesystem-test-{uuid.uuid4().hex[:8]}",
|
||||
"metadata": {},
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert create_resp.status_code == 200, create_resp.text
|
||||
payload = create_resp.json()
|
||||
thread_id = payload.get("thread_id") or payload.get("id")
|
||||
assert thread_id, f"Create thread response missing thread identifier: {payload}"
|
||||
return thread_id
|
||||
|
||||
|
||||
async def _upload_attachment_file(
|
||||
test_client,
|
||||
thread_id: str,
|
||||
headers: dict[str, str],
|
||||
file_name: str,
|
||||
content: str,
|
||||
) -> str:
|
||||
response = await test_client.post(
|
||||
f"/api/chat/thread/{thread_id}/attachments",
|
||||
files={"file": (file_name, content.encode("utf-8"), "text/plain")},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
payload = response.json()
|
||||
path = payload.get("path") or payload.get("file_path")
|
||||
assert isinstance(path, str) and path
|
||||
return path
|
||||
|
||||
|
||||
async def test_viewer_tree_requires_authentication(test_client):
|
||||
response = await test_client.get("/api/viewer/filesystem/tree", params={"thread_id": "x", "path": "/"})
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
async def test_viewer_tree_root_lists_user_data_namespace(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
entries = response.json().get("entries", [])
|
||||
paths = {entry.get("path") for entry in entries}
|
||||
assert "/home/gem/user-data/" in paths
|
||||
|
||||
|
||||
async def test_viewer_tree_root_does_not_require_sandbox_listing(test_client, standard_user, monkeypatch):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
class _FailingSandbox:
|
||||
def ls(self, path):
|
||||
raise AssertionError(f"sandbox ls should not be used for root path: {path}")
|
||||
|
||||
class _EmptyBackend:
|
||||
def has_entries(self):
|
||||
return False
|
||||
|
||||
service_module = importlib.import_module("yuxi.services.viewer_filesystem_service")
|
||||
|
||||
async def _fake_resolve_viewer_state(**kwargs):
|
||||
return _FailingSandbox(), _EmptyBackend(), []
|
||||
|
||||
monkeypatch.setattr(service_module, "_resolve_viewer_state", _fake_resolve_viewer_state)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
entries = response.json().get("entries", [])
|
||||
paths = {entry.get("path") for entry in entries}
|
||||
assert "/home/gem/user-data/" in paths
|
||||
|
||||
|
||||
async def test_viewer_tree_user_data_uses_local_thread_directory(test_client, standard_user, monkeypatch):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "viewer_tree_demo.txt"
|
||||
actual_path.write_text("viewer tree", encoding="utf-8")
|
||||
|
||||
class _FailingSandbox:
|
||||
def ls(self, path):
|
||||
raise AssertionError(f"sandbox ls should not be used for user-data path: {path}")
|
||||
|
||||
class _EmptyBackend:
|
||||
def has_entries(self):
|
||||
return False
|
||||
|
||||
service_module = importlib.import_module("yuxi.services.viewer_filesystem_service")
|
||||
|
||||
async def _fake_resolve_viewer_state(**kwargs):
|
||||
return _FailingSandbox(), _EmptyBackend(), []
|
||||
|
||||
monkeypatch.setattr(service_module, "_resolve_viewer_state", _fake_resolve_viewer_state)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
entries = response.json().get("entries", [])
|
||||
paths = {entry.get("path") for entry in entries}
|
||||
assert "/home/gem/user-data/workspace/" in paths
|
||||
assert "/home/gem/user-data/uploads/" in paths
|
||||
assert "/home/gem/user-data/outputs/" in paths
|
||||
assert all(str(path).startswith("/home/gem/user-data") for path in paths)
|
||||
|
||||
|
||||
async def test_viewer_tree_user_data_root_keeps_thread_root_files_visible(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
root_file = sandbox_user_data_dir(thread_id) / "root-note.txt"
|
||||
root_file.write_text("visible at root", encoding="utf-8")
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
paths = {entry.get("path") for entry in response.json().get("entries", [])}
|
||||
assert "/home/gem/user-data/root-note.txt" in paths
|
||||
assert "/home/gem/user-data/workspace/" in paths
|
||||
|
||||
|
||||
async def test_workspace_is_shared_across_same_user_threads_but_uploads_remain_thread_local(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
other_thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
shared_path = sandbox_workspace_dir(thread_id, uid) / "shared_across_threads.txt"
|
||||
shared_path.write_text("shared workspace", encoding="utf-8")
|
||||
|
||||
upload_path = await _upload_attachment_file(test_client, thread_id, headers, "thread-local.txt", "private upload\n")
|
||||
|
||||
workspace_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": other_thread_id, "path": "/home/gem/user-data/workspace"},
|
||||
headers=headers,
|
||||
)
|
||||
assert workspace_response.status_code == 200, workspace_response.text
|
||||
workspace_paths = {entry.get("path") for entry in workspace_response.json().get("entries", [])}
|
||||
assert "/home/gem/user-data/workspace/shared_across_threads.txt" in workspace_paths
|
||||
|
||||
shared_file_response = await test_client.get(
|
||||
f"/api/chat/thread/{other_thread_id}/artifacts/home/gem/user-data/workspace/shared_across_threads.txt",
|
||||
headers=headers,
|
||||
)
|
||||
assert shared_file_response.status_code == 200, shared_file_response.text
|
||||
assert shared_file_response.text == "shared workspace"
|
||||
|
||||
upload_response = await test_client.get(
|
||||
f"/api/chat/thread/{other_thread_id}/artifacts/{upload_path.lstrip('/')}",
|
||||
headers=headers,
|
||||
)
|
||||
assert upload_response.status_code == 404, upload_response.text
|
||||
|
||||
|
||||
async def test_viewer_file_returns_raw_content_without_line_numbers(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
file_path = await _upload_attachment_file(test_client, thread_id, headers, "viewer_demo.txt", "alpha\nbeta\n")
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert "alpha" in response.json()["content"]
|
||||
assert "beta" in response.json()["content"]
|
||||
assert response.json()["preview_type"] == "markdown"
|
||||
assert response.json()["supported"] is True
|
||||
|
||||
|
||||
async def test_viewer_file_returns_unsupported_for_binary_payload(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "binary.bin"
|
||||
actual_path.write_bytes(b"\x00\x01\x02\x03binary")
|
||||
file_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["content"] is None
|
||||
assert payload["preview_type"] == "unsupported"
|
||||
assert payload["supported"] is False
|
||||
|
||||
|
||||
async def test_viewer_file_returns_image_preview_metadata(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "demo.png"
|
||||
actual_path.write_bytes(
|
||||
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06\x00\x00\x00\x1f\x15\xc4\x89"
|
||||
)
|
||||
file_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["content"] is None
|
||||
assert payload["preview_type"] == "image"
|
||||
assert payload["supported"] is True
|
||||
|
||||
|
||||
async def test_viewer_file_returns_pdf_preview_metadata(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "demo.pdf"
|
||||
actual_path.write_bytes(b"%PDF-1.4\n1 0 obj\n<<>>\nendobj\ntrailer\n<<>>\n%%EOF")
|
||||
file_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["content"] is None
|
||||
assert payload["preview_type"] == "pdf"
|
||||
assert payload["supported"] is True
|
||||
|
||||
|
||||
async def test_viewer_download_returns_attachment_response(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
file_path = await _upload_attachment_file(test_client, thread_id, headers, "download_demo.txt", "download-me\n")
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/download",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
content_disposition = response.headers.get("content-disposition", "")
|
||||
assert "attachment;" in content_disposition
|
||||
assert "download_demo" in content_disposition
|
||||
assert "download-me" in response.text
|
||||
|
||||
|
||||
async def test_viewer_download_returns_full_file_for_large_user_data_content(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
large_content = "0123456789abcdef" * 4096
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "large_download.txt"
|
||||
actual_path.write_text(large_content, encoding="utf-8")
|
||||
file_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/download",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.content == large_content.encode("utf-8")
|
||||
|
||||
|
||||
async def test_viewer_delete_removes_user_data_file(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "delete_me.txt"
|
||||
actual_path.write_text("delete me", encoding="utf-8")
|
||||
file_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
delete_response = await test_client.delete(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
assert delete_response.json()["success"] is True
|
||||
assert not actual_path.exists()
|
||||
|
||||
tree_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data/workspace"},
|
||||
headers=headers,
|
||||
)
|
||||
assert tree_response.status_code == 200, tree_response.text
|
||||
paths = {entry.get("path") for entry in tree_response.json().get("entries", [])}
|
||||
assert file_path not in paths
|
||||
|
||||
|
||||
async def test_viewer_delete_removes_empty_user_data_directory(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "empty-folder"
|
||||
actual_path.mkdir()
|
||||
dir_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
delete_response = await test_client.delete(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": dir_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
assert delete_response.json()["success"] is True
|
||||
assert not actual_path.exists()
|
||||
|
||||
tree_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data/workspace"},
|
||||
headers=headers,
|
||||
)
|
||||
assert tree_response.status_code == 200, tree_response.text
|
||||
paths = {entry.get("path") for entry in tree_response.json().get("entries", [])}
|
||||
assert f"{dir_path}/" not in paths
|
||||
|
||||
|
||||
async def test_viewer_delete_recursively_removes_user_data_directory(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
actual_path = sandbox_workspace_dir(thread_id, uid) / "nested-folder"
|
||||
nested_dir = actual_path / "child"
|
||||
nested_dir.mkdir(parents=True)
|
||||
nested_file = nested_dir / "notes.txt"
|
||||
nested_file.write_text("remove recursively", encoding="utf-8")
|
||||
dir_path = virtual_path_for_thread_file(thread_id, actual_path, uid=uid)
|
||||
|
||||
delete_response = await test_client.delete(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": dir_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert delete_response.status_code == 200, delete_response.text
|
||||
assert delete_response.json()["success"] is True
|
||||
assert not actual_path.exists()
|
||||
assert not nested_file.exists()
|
||||
|
||||
tree_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data/workspace"},
|
||||
headers=headers,
|
||||
)
|
||||
assert tree_response.status_code == 200, tree_response.text
|
||||
paths = {entry.get("path") for entry in tree_response.json().get("entries", [])}
|
||||
assert f"{dir_path}/" not in paths
|
||||
|
||||
|
||||
async def test_viewer_delete_rejects_readonly_namespace_directory(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
response = await test_client.delete(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/skills"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"] == "当前路径不支持删除"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"protected_path",
|
||||
[
|
||||
"/home/gem/user-data",
|
||||
"/home/gem/user-data/workspace",
|
||||
"/home/gem/user-data/uploads",
|
||||
"/home/gem/user-data/outputs",
|
||||
],
|
||||
)
|
||||
async def test_viewer_delete_rejects_protected_user_data_root_directories(
|
||||
test_client, standard_user, protected_path: str
|
||||
):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
|
||||
response = await test_client.delete(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": protected_path},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"] == "当前目录不允许删除"
|
||||
|
||||
|
||||
async def test_delete_viewer_file_rejects_user_data_root_without_removing(monkeypatch):
|
||||
from yuxi.services import viewer_filesystem_service as service_module
|
||||
|
||||
async def _fake_resolve_viewer_state(**kwargs):
|
||||
return object(), object(), []
|
||||
|
||||
def _fail_rmtree(path):
|
||||
raise AssertionError(f"unexpected root removal: {path}")
|
||||
|
||||
monkeypatch.setattr(service_module, "_resolve_viewer_state", _fake_resolve_viewer_state)
|
||||
monkeypatch.setattr(service_module.shutil, "rmtree", _fail_rmtree)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await service_module.delete_viewer_file(
|
||||
thread_id="thread-1",
|
||||
path="/home/gem/user-data",
|
||||
current_user=SimpleNamespace(uid="user-1"),
|
||||
db=None,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.detail == "当前目录不允许删除"
|
||||
|
||||
|
||||
async def test_viewer_create_directory_adds_workspace_folder(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/directory",
|
||||
json={
|
||||
"thread_id": thread_id,
|
||||
"parent_path": "/home/gem/user-data/workspace",
|
||||
"name": "created-folder",
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert payload["entry"]["path"] == "/home/gem/user-data/workspace/created-folder/"
|
||||
assert (sandbox_workspace_dir(thread_id, uid) / "created-folder").is_dir()
|
||||
|
||||
tree_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data/workspace"},
|
||||
headers=headers,
|
||||
)
|
||||
assert tree_response.status_code == 200, tree_response.text
|
||||
paths = {entry.get("path") for entry in tree_response.json().get("entries", [])}
|
||||
assert "/home/gem/user-data/workspace/created-folder/" in paths
|
||||
|
||||
|
||||
async def test_viewer_upload_file_writes_to_workspace_subdirectory(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
target_dir = sandbox_workspace_dir(thread_id, uid) / "upload-target"
|
||||
target_dir.mkdir()
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/upload",
|
||||
data={"thread_id": thread_id, "parent_path": "/home/gem/user-data/workspace/upload-target"},
|
||||
files={"files": ("uploaded.txt", b"uploaded from viewer\n", "text/plain")},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
payload = response.json()
|
||||
assert payload["success"] is True
|
||||
assert payload["entries"][0]["path"] == "/home/gem/user-data/workspace/upload-target/uploaded.txt"
|
||||
|
||||
file_response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/user-data/workspace/upload-target/uploaded.txt"},
|
||||
headers=headers,
|
||||
)
|
||||
assert file_response.status_code == 200, file_response.text
|
||||
assert file_response.json()["content"] == "uploaded from viewer\n"
|
||||
|
||||
|
||||
async def test_viewer_create_directory_rejects_conflict(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
(sandbox_workspace_dir(thread_id, uid) / "conflict").mkdir()
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/directory",
|
||||
json={
|
||||
"thread_id": thread_id,
|
||||
"parent_path": "/home/gem/user-data/workspace",
|
||||
"name": "conflict",
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"] == "同名文件或文件夹已存在"
|
||||
|
||||
|
||||
async def test_viewer_upload_file_rejects_conflict_without_overwrite(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
existing_file = sandbox_workspace_dir(thread_id, uid) / "existing.txt"
|
||||
existing_file.write_text("keep me\n", encoding="utf-8")
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/upload",
|
||||
data={"thread_id": thread_id, "parent_path": "/home/gem/user-data/workspace"},
|
||||
files={"files": ("existing.txt", b"replace me\n", "text/plain")},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json()["detail"] == "同名文件或文件夹已存在"
|
||||
assert existing_file.read_text(encoding="utf-8") == "keep me\n"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parent_path",
|
||||
[
|
||||
"/home/gem/skills",
|
||||
"/home/gem/user-data/uploads",
|
||||
"/home/gem/user-data/outputs",
|
||||
],
|
||||
)
|
||||
async def test_viewer_write_rejects_non_workspace_paths(test_client, standard_user, parent_path: str):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
directory_response = await test_client.post(
|
||||
"/api/viewer/filesystem/directory",
|
||||
json={"thread_id": thread_id, "parent_path": parent_path, "name": "blocked"},
|
||||
headers=headers,
|
||||
)
|
||||
assert directory_response.status_code == 400, directory_response.text
|
||||
assert directory_response.json()["detail"] == "当前路径不支持写入"
|
||||
|
||||
upload_response = await test_client.post(
|
||||
"/api/viewer/filesystem/upload",
|
||||
data={"thread_id": thread_id, "parent_path": parent_path},
|
||||
files={"files": ("blocked.txt", b"blocked", "text/plain")},
|
||||
headers=headers,
|
||||
)
|
||||
assert upload_response.status_code == 400, upload_response.text
|
||||
assert upload_response.json()["detail"] == "当前路径不支持写入"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("folder_name", ["", "../escape", "nested/folder"])
|
||||
async def test_viewer_create_directory_rejects_invalid_names(test_client, standard_user, folder_name: str):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/directory",
|
||||
json={
|
||||
"thread_id": thread_id,
|
||||
"parent_path": "/home/gem/user-data/workspace",
|
||||
"name": folder_name,
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 422, response.text
|
||||
|
||||
|
||||
async def test_viewer_tree_root_hides_kbs_namespace(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/"},
|
||||
headers=headers,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
entries = response.json().get("entries", [])
|
||||
paths = {entry.get("path") for entry in entries}
|
||||
assert "/home/gem/user-data/" in paths
|
||||
assert "/home/gem/kbs/" not in paths
|
||||
|
||||
|
||||
async def test_viewer_rejects_kbs_namespace(test_client, standard_user):
|
||||
headers = standard_user["headers"]
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
tree_response = await test_client.get(
|
||||
"/api/viewer/filesystem/tree",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/kbs"},
|
||||
headers=headers,
|
||||
)
|
||||
assert tree_response.status_code == 400, tree_response.text
|
||||
|
||||
file_response = await test_client.get(
|
||||
"/api/viewer/filesystem/file",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/kbs/demo.md"},
|
||||
headers=headers,
|
||||
)
|
||||
assert file_response.status_code == 400, file_response.text
|
||||
|
||||
download_response = await test_client.get(
|
||||
"/api/viewer/filesystem/download",
|
||||
params={"thread_id": thread_id, "path": "/home/gem/kbs/demo.md"},
|
||||
headers=headers,
|
||||
)
|
||||
assert download_response.status_code == 400, download_response.text
|
||||
@@ -0,0 +1,78 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from yuxi.agents.backends.sandbox import ensure_thread_dirs, sandbox_workspace_dir
|
||||
|
||||
pytestmark = [pytest.mark.asyncio, pytest.mark.integration]
|
||||
|
||||
|
||||
async def _create_thread_for_user(test_client, headers: dict[str, str]) -> str:
|
||||
agent_resp = await test_client.get("/api/agent/default", headers=headers)
|
||||
assert agent_resp.status_code == 200, agent_resp.text
|
||||
agent = agent_resp.json().get("agent") or {}
|
||||
agent_id = agent.get("slug") or agent.get("id")
|
||||
if not agent_id:
|
||||
pytest.skip("Default agent payload missing id field.")
|
||||
|
||||
create_resp = await test_client.post(
|
||||
"/api/chat/thread",
|
||||
json={
|
||||
"agent_id": agent_id,
|
||||
"title": f"viewer-security-test-{uuid.uuid4().hex[:8]}",
|
||||
"metadata": {},
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
assert create_resp.status_code == 200, create_resp.text
|
||||
payload = create_resp.json()
|
||||
thread_id = payload.get("thread_id") or payload.get("id")
|
||||
assert thread_id
|
||||
return thread_id
|
||||
|
||||
|
||||
async def test_viewer_download_blocks_workspace_symlink_escape(test_client, standard_user, tmp_path):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
workspace_dir = sandbox_workspace_dir(thread_id, uid)
|
||||
outside_file = tmp_path / "outside.txt"
|
||||
outside_file.write_text("outside", encoding="utf-8")
|
||||
symlink_path = workspace_dir / "escape.txt"
|
||||
symlink_path.symlink_to(outside_file)
|
||||
file_path = "/home/gem/user-data/workspace/escape.txt"
|
||||
|
||||
response = await test_client.get(
|
||||
"/api/viewer/filesystem/download",
|
||||
params={"thread_id": thread_id, "path": file_path},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
|
||||
|
||||
async def test_viewer_upload_blocks_workspace_symlink_escape(test_client, standard_user, tmp_path):
|
||||
headers = standard_user["headers"]
|
||||
uid = str(standard_user["user"]["uid"])
|
||||
thread_id = await _create_thread_for_user(test_client, headers)
|
||||
|
||||
ensure_thread_dirs(thread_id, uid)
|
||||
workspace_dir = sandbox_workspace_dir(thread_id, uid)
|
||||
outside_dir = tmp_path / "outside-dir"
|
||||
outside_dir.mkdir()
|
||||
symlink_path = workspace_dir / "escape-dir"
|
||||
symlink_path.symlink_to(outside_dir)
|
||||
parent_path = "/home/gem/user-data/workspace/escape-dir"
|
||||
|
||||
response = await test_client.post(
|
||||
"/api/viewer/filesystem/upload",
|
||||
data={"thread_id": thread_id, "parent_path": parent_path},
|
||||
files={"files": ("escape.txt", b"outside", "text/plain")},
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert not (outside_dir / "escape.txt").exists()
|
||||
Reference in New Issue
Block a user