mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 19:22:00 +00:00
ThearchyHelios
463 lines
23 KiB
Python
463 lines
23 KiB
Python
"""开登录后使用者的问数据（D-c）：能发起什么样的运行、查询结果在工具边界遮、运行的 WebSocket 按运行判、数据源名单。
|
|
|
|
- 图：只许输入、模型、Agent、报告、成果、知识检索、数据整形、条件分支、口径卡、合并查询；Agent 只许数据源工具和只读的
|
|
内置工具（app/auth/chat_policy.py）。调用工具、沙箱代码、自定义工具、MCP、网页、HTTP、文件一律 403；
|
|
- 使用者发起的问数据运行一定带 MASK_RESULTS_OPTION：真的执行一次（演示模型调数据源查询），事件、快照、实时流里都没有
|
|
遮罩列的原值；
|
|
- WebSocket：没登录的握手前关掉；使用者只能连自己能读、查询结果在工具边界遮过的运行（握手时判，1008 拒）；
|
|
- GET /api/datasources/names：开了问数据的使用者能读，只有名字、说明、种类，只列启用着的。
|
|
|
|
数据源是现造的 SQLite 文件，姓名、手机号都是虚构的。
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import sqlite3
|
|
import uuid
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
from sqlalchemy import delete, select
|
|
|
|
from app.auth.chat_policy import chat_graph_problem
|
|
from app.data import introspect
|
|
from app.db.base import SessionLocal
|
|
from app.db.models import Artifact, DataSource, Run, RunEvent
|
|
from app.engine.runner import run_manager
|
|
from app.tools.datasource import MASK_RESULTS_OPTION
|
|
from tests.fakes import auth_env
|
|
|
|
PHONE = "13900000001"
|
|
MASKED = "已遮罩"
|
|
|
|
|
|
def node(nid, ntype, **config):
|
|
return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}}
|
|
|
|
|
|
def chain(*nodes):
|
|
return {"nodes": list(nodes), "edges": [{"source": a["id"], "target": b["id"]} for a, b in zip(nodes, nodes[1:])]}
|
|
|
|
|
|
def chat_graph(*tools: str, middle: dict | None = None) -> dict:
|
|
return chain(node("start", "input", fields=[{"name": "question"}]),
|
|
middle or node("ask", "agent", prompt="{{ input.question }}", tools=list(tools), max_steps=4),
|
|
node("out", "output", fields=[{"name": "回答", "value": "{{ nodes.ask.text }}"}]))
|
|
|
|
|
|
# ---------------------------------------------------------------- 图的白名单（纯函数）
|
|
|
|
|
|
@pytest.mark.parametrize("graph", [
|
|
chat_graph("db_query__crm", "db_schema__crm", "calculator", "current_time", "json_query", "kb_search",
|
|
"memory_recall"),
|
|
chain(node("start", "input"), node("kb", "retrieve", query="x"), node("t", "transform", expression="1"),
|
|
node("out", "output")),
|
|
{"nodes": "坏的"}, # 不成形的图交给 launch_run 照常报校验问题
|
|
])
|
|
def test_graphs_a_user_may_run_in_chat(graph):
|
|
assert chat_graph_problem(graph) is None
|
|
|
|
|
|
@pytest.mark.parametrize("graph, said", [
|
|
(chat_graph(middle=node("ask", "code", code="print(1)")), "沙箱代码"),
|
|
(chat_graph(middle=node("ask", "tool", tool="calculator")), "调用工具"),
|
|
(chat_graph(middle=node("ask", "supervisor", agents=[])), "多 Agent 协作"),
|
|
(chat_graph(middle=node("ask", "subgraph", workflow_id="w")), "子工作流"),
|
|
(chat_graph(middle=node("ask", "human", mode="approve")), "人工审批"),
|
|
(chat_graph(middle=node("ask", "memory", action="write")), "长期记忆"),
|
|
(chat_graph("http_request"), "http_request"),
|
|
(chat_graph("web_fetch"), "web_fetch"),
|
|
(chat_graph("python_exec"), "python_exec"),
|
|
(chat_graph("file_read"), "file_read"),
|
|
(chat_graph("memory_write"), "memory_write"),
|
|
(chat_graph("mcp:fs/read_file"), "mcp:fs/read_file"),
|
|
(chat_graph("my_custom_http"), "my_custom_http"),
|
|
(chat_graph("{{ vars.tool }}"), "模板"),
|
|
])
|
|
def test_graphs_a_user_may_not_run_in_chat(graph, said):
|
|
problem = chat_graph_problem(graph)
|
|
assert problem and problem.startswith("没有权限") and said in problem, problem
|
|
|
|
|
|
@pytest.mark.parametrize("graph", [
|
|
chat_graph(middle=node("ask", "agent", prompt="x", tools=["memory_recall"], memory_scope="builders")),
|
|
chat_graph(middle=node("kb", "retrieve", query="x", collection="hr")),
|
|
{**chat_graph("kb_search"), "defaults": {"collection": "hr"}},
|
|
{**chat_graph("memory_recall"), "defaults": {"memory_scope": " "}},
|
|
])
|
|
def test_users_may_not_pick_memory_or_knowledge_scope(graph):
|
|
"""记忆、知识库的内容不经遮罩：节点上、图级 defaults 里都不许指定范围，用管理员设好的默认值。"""
|
|
problem = chat_graph_problem(graph)
|
|
assert problem and "记忆范围和知识库" in problem
|
|
|
|
|
|
def test_tools_from_graph_defaults_count_too():
|
|
"""Agent 上没写工具时运行时用图级默认值里的：白名单同样按它查（engine/launch.tool_bindings）。"""
|
|
graph = chain(node("start", "input"), node("ask", "agent", prompt="x"), node("out", "output"))
|
|
graph["defaults"] = {"tools": ["web_search"]}
|
|
problem = chat_graph_problem(graph)
|
|
assert problem and "web_search" in problem
|
|
|
|
|
|
# ---------------------------------------------------------------- 接口
|
|
|
|
|
|
async def _masked_source(tmp_path) -> str:
|
|
path = tmp_path / f"{uuid.uuid4().hex[:8]}.db"
|
|
db = sqlite3.connect(path)
|
|
db.executescript("CREATE TABLE customers (id INTEGER PRIMARY KEY, name TEXT, phone TEXT, city TEXT);"
|
|
f"INSERT INTO customers VALUES (1, '张三', '{PHONE}', '杭州'), (2, '李四', '13900000002', '宁波');")
|
|
db.commit()
|
|
db.close()
|
|
name = f"crm_{uuid.uuid4().hex[:6]}"
|
|
async with SessionLocal() as session:
|
|
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=True, enabled=True,
|
|
description="客户", options={"mask_columns": "phone"})
|
|
session.add(row)
|
|
await session.flush()
|
|
row.schema_cache = await introspect.introspect(row)
|
|
session.add(DataSource(name=f"off_{uuid.uuid4().hex[:6]}", kind="sqlite", database=str(path), enabled=False))
|
|
await session.commit()
|
|
return name
|
|
|
|
|
|
@pytest.fixture
|
|
async def world(tmp_path):
|
|
await auth_env.reset()
|
|
await run_manager.setup()
|
|
admin = await auth_env.turn_on_local()
|
|
await auth_env.make_user("uma", "user")
|
|
await auth_env.make_user("bob", "builder")
|
|
source = await _masked_source(tmp_path)
|
|
uma = await auth_env.login("uma")
|
|
bob = await auth_env.login("bob")
|
|
try:
|
|
yield {"admin": admin, "uma": uma, "bob": bob, "source": source}
|
|
finally:
|
|
await run_manager.shutdown()
|
|
async with SessionLocal() as s:
|
|
await s.execute(delete(Run).where(Run.started_by.in_(["uma", "bob"])))
|
|
await s.execute(delete(DataSource))
|
|
await s.commit()
|
|
await auth_env.reset()
|
|
|
|
|
|
async def _allow_chat(world) -> None:
|
|
r = await world["admin"].put("/api/auth/settings", json={"user_can_chat": True})
|
|
assert r.status_code == 200, r.text
|
|
|
|
|
|
async def _finish(run_id: str) -> Run:
|
|
for _ in range(400):
|
|
async with SessionLocal() as session:
|
|
row = await session.get(Run, run_id)
|
|
if row.status in ("succeeded", "failed", "interrupted", "cancelled"):
|
|
await run_manager.wait_idle(run_id)
|
|
return row
|
|
await asyncio.sleep(0.05)
|
|
raise AssertionError("运行没有结束")
|
|
|
|
|
|
def _script(monkeypatch, tool: str) -> None:
|
|
"""演示模型：第一轮查遮罩列，看到结果后作答。"""
|
|
from app.providers import mock_model
|
|
|
|
def _decide(self, messages):
|
|
if any(isinstance(m, ToolMessage) for m in messages):
|
|
return AIMessage(content="杭州有一位客户。")
|
|
return AIMessage(content="", tool_calls=[{"name": tool, "args": {"sql": "SELECT name, phone, city FROM customers"},
|
|
"id": "call_1"}])
|
|
|
|
monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
|
|
|
|
|
|
async def test_a_user_can_only_launch_whitelisted_chat_graphs(world):
|
|
await _allow_chat(world)
|
|
uma = world["uma"]
|
|
r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph("http_request"), "input": {}})
|
|
assert r.status_code == 403 and "http_request" in r.json()["detail"]
|
|
r = await uma.post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
|
|
assert r.status_code == 403 and "沙箱代码" in r.json()["detail"]
|
|
r = await uma.post("/api/runs", json={"origin": "chat", "run_class": "formal", "workflow_id": "x"})
|
|
assert r.status_code == 403, "正式运行只经应用入口发起"
|
|
# 搭建者照旧不受这份白名单约束（发起前检查另说：这里只看不是 403）
|
|
r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(middle=node("ask", "code"))})
|
|
assert r.status_code != 403
|
|
|
|
|
|
async def test_a_users_chat_run_is_masked_end_to_end(world, monkeypatch):
|
|
"""真的执行一次：使用者发起、Agent 查了遮罩列。运行选项带着遮罩开关，事件、查询快照、实时流里都没有原值。"""
|
|
await _allow_chat(world)
|
|
tool = f"db_query__{world['source']}"
|
|
_script(monkeypatch, tool)
|
|
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": chat_graph(tool),
|
|
"input": {"question": "杭州有几位客户"}})
|
|
assert r.status_code == 201, r.text
|
|
run = await _finish(r.json()["id"])
|
|
assert run.status == "succeeded", run.error
|
|
assert run.options.get(MASK_RESULTS_OPTION) is True and run.started_by == "uma" and run.origin_kind == "chat"
|
|
async with SessionLocal() as s:
|
|
events = list((await s.execute(select(RunEvent).where(RunEvent.run_id == run.id))).scalars())
|
|
snaps = list((await s.execute(select(Artifact).where(Artifact.run_id == run.id,
|
|
Artifact.kind == "query_snapshot"))).scalars())
|
|
dumped = json.dumps([e.data for e in events], ensure_ascii=False)
|
|
assert "tool.end" in {e.type for e in events}
|
|
assert PHONE not in dumped and "张三" in dumped, "遮罩列换掉，别的列照旧"
|
|
assert snaps, "查询照常落快照"
|
|
from app.core.artifact_store import load
|
|
|
|
for snap in snaps:
|
|
assert PHONE not in json.dumps(load(snap.id), ensure_ascii=False)
|
|
|
|
# 使用者能读自己这次运行的记录、事件，也能连实时流；实时流里同样没有原值
|
|
assert (await world["uma"].get(f"/api/runs/{run.id}/events")).status_code == 200
|
|
frames = await _ws(f"/api/runs/{run.id}/stream", _cookie(world["uma"]))
|
|
assert frames[0]["type"] == "websocket.accept"
|
|
texts = [f.get("text") or "" for f in frames if f["type"] == "websocket.send"]
|
|
assert texts and '"stream.end"' in texts[-1] and PHONE not in "".join(texts)
|
|
|
|
|
|
async def test_a_users_chat_run_cannot_write_a_writable_source(world, monkeypatch, tmp_path):
|
|
"""可写源 + Agent 自己写了 approval=never：写语句照样被守卫拒掉、不执行（只读不交给审批），表里的行一条没少。
|
|
搭建者发起的同一张图照旧能写——只读只给 user 角色的问数据运行。"""
|
|
from app.providers import mock_model
|
|
from app.tools.datasource import READ_ONLY_OPTION
|
|
|
|
path = tmp_path / "ledger.db"
|
|
db = sqlite3.connect(path)
|
|
db.executescript("CREATE TABLE ledger (id INTEGER PRIMARY KEY, amount INTEGER);"
|
|
"INSERT INTO ledger VALUES (1, 10), (2, 20);")
|
|
db.commit()
|
|
db.close()
|
|
name = f"ledger_{uuid.uuid4().hex[:6]}"
|
|
async with SessionLocal() as session:
|
|
row = DataSource(name=name, kind="sqlite", database=str(path), readonly=False, enabled=True)
|
|
session.add(row)
|
|
await session.flush()
|
|
row.schema_cache = await introspect.introspect(row)
|
|
await session.commit()
|
|
replies: list[str] = []
|
|
|
|
def _decide(self, messages):
|
|
if seen := [m for m in messages if isinstance(m, ToolMessage)]:
|
|
replies.append(str(seen[-1].content))
|
|
return AIMessage(content="好了。")
|
|
return AIMessage(content="", tool_calls=[{"name": f"db_query__{name}", "args": {"sql": "DELETE FROM ledger"},
|
|
"id": "call_1"}])
|
|
|
|
monkeypatch.setattr(mock_model.MockChatModel, "_decide", _decide)
|
|
graph = chat_graph(middle=node("ask", "agent", prompt="{{ input.question }}", tools=[f"db_query__{name}"],
|
|
max_steps=4, approval="never"))
|
|
await _allow_chat(world)
|
|
r = await world["uma"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
|
|
assert r.status_code == 201, r.text
|
|
run = await _finish(r.json()["id"])
|
|
assert run.status == "succeeded", run.error
|
|
assert run.options.get(READ_ONLY_OPTION) is True
|
|
assert replies and replies[-1].startswith("SQL 被拒绝"), replies
|
|
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 2
|
|
|
|
replies.clear()
|
|
r = await world["bob"].post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "清空"}})
|
|
assert r.status_code == 201, r.text
|
|
run = await _finish(r.json()["id"])
|
|
assert run.status == "succeeded", run.error
|
|
assert not (run.options or {}).get(READ_ONLY_OPTION)
|
|
assert sqlite3.connect(path).execute("SELECT COUNT(*) FROM ledger").fetchone()[0] == 0, "搭建者照旧能写"
|
|
|
|
|
|
async def test_a_users_chat_run_keeps_to_its_own_memory_scope(world):
|
|
"""请求体里带记忆范围、知识库的拒；不带的，记忆是这位使用者自己的范围，知识库是设置里的默认值。"""
|
|
await _allow_chat(world)
|
|
uma = world["uma"]
|
|
graph = chat_graph("memory_recall")
|
|
for extra in ({"memory_scope": "default"}, {"collection": "hr"}):
|
|
r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {}, **extra})
|
|
assert r.status_code == 403 and "记忆范围和知识库" in r.json()["detail"]
|
|
r = await uma.post("/api/runs", json={"origin": "chat", "graph": graph, "input": {"question": "x"}})
|
|
assert r.status_code == 201, r.text
|
|
body = r.json()
|
|
assert body["memory_scope"] == "chat:uma" and body["collection"] == "default"
|
|
await _finish(body["id"])
|
|
|
|
|
|
# ---------------------------------------------------------------- 问数据的生成流
|
|
|
|
|
|
class _Scripted:
|
|
"""第一轮按节点手册的标准写法搭（tool 节点查库），收到自查请求时改成 Agent 查库。记下 system 和每一轮的请求。"""
|
|
|
|
def __init__(self, source: str) -> None:
|
|
self.source = source
|
|
self.systems: list[str] = []
|
|
self.humans: list[str] = []
|
|
|
|
def _lines(self, repair: bool) -> list[str]:
|
|
if repair:
|
|
ops = [{"op": "remove_node", "id": "fetch"},
|
|
{"op": "add_node", "node": {"id": "ask", "type": "agent", "label": "查库",
|
|
"config": {"prompt": "{{ input.question }}",
|
|
"tools": [f"db_query__{self.source}"]}}},
|
|
{"op": "add_edge", "edge": {"source": "start", "target": "ask"}},
|
|
{"op": "add_edge", "edge": {"source": "ask", "target": "out"}},
|
|
{"op": "done", "explanation": "改成 Agent 查库"}]
|
|
else:
|
|
ops = [{"op": "plan", "summary": "查客户"},
|
|
{"op": "add_node", "node": {"id": "start", "type": "input", "label": "输入",
|
|
"config": {"fields": [{"name": "question"}]}}},
|
|
{"op": "add_node", "node": {"id": "fetch", "type": "tool", "label": "查库",
|
|
"config": {"tool": f"db_query__{self.source}",
|
|
"args": {"sql": "SELECT city FROM customers"}}}},
|
|
{"op": "add_node", "node": {"id": "out", "type": "output", "label": "成果",
|
|
"config": {"fields": [{"name": "回答", "value": "{{ last_message }}"}]}}},
|
|
{"op": "add_edge", "edge": {"source": "start", "target": "fetch"}},
|
|
{"op": "add_edge", "edge": {"source": "fetch", "target": "out"}},
|
|
{"op": "done", "explanation": "查好了", "run": True}]
|
|
return [json.dumps(o, ensure_ascii=False) for o in ops]
|
|
|
|
async def astream(self, messages):
|
|
from langchain_core.messages import AIMessageChunk
|
|
|
|
self.systems.append(next(text for role, text in messages if role == "system"))
|
|
human = next(text for role, text in messages if role == "human")
|
|
self.humans.append(human)
|
|
for line in self._lines("上线前自查" in human):
|
|
yield AIMessageChunk(content=line + "\n")
|
|
|
|
|
|
async def _generate(client, monkeypatch, model: _Scripted) -> list[dict]:
|
|
import app.api.copilot as copilot
|
|
|
|
async def _model(*_a, **_k):
|
|
return model, "mock-fast"
|
|
|
|
monkeypatch.setattr(copilot, "get_chat_model", _model)
|
|
events: list[dict] = []
|
|
async with client.stream("POST", "/api/copilot/generate-stream", json={"instruction": "杭州有几位客户"}) as r:
|
|
assert r.status_code == 200
|
|
async for line in r.aiter_lines():
|
|
if line.startswith("data: "):
|
|
events.append(json.loads(line[6:]))
|
|
return events
|
|
|
|
|
|
async def test_generation_for_a_user_keeps_to_the_chat_policy(world, monkeypatch):
|
|
"""使用者问数据时，提示词先说清只能 Agent 查库；照标准写法搭出的 tool 节点图按 chat_policy 交回去改，改好才交付、
|
|
自动运行。搭建者照旧按节点手册来，tool 节点不算问题。"""
|
|
await _allow_chat(world)
|
|
model = _Scripted(world["source"])
|
|
events = await _generate(world["uma"], monkeypatch, model)
|
|
assert "这位用户是「使用者」" in model.systems[0] and "- http_request" not in model.systems[0]
|
|
assert len(model.humans) == 2 and "调用工具" in model.humans[1], "tool 节点按 chat_policy 交回去改"
|
|
check = [e for e in events if e.get("op") == "check"][-1]
|
|
final = next(e for e in events if e.get("op") == "final")
|
|
assert check["status"] == "passed" and final["autorun"] is True
|
|
assert {n["type"] for n in final["graph"]["nodes"]} == {"input", "agent", "output"}
|
|
assert chat_graph_problem(final["graph"]) is None
|
|
|
|
model = _Scripted(world["source"])
|
|
events = await _generate(world["bob"], monkeypatch, model)
|
|
assert "这位用户是「使用者」" not in model.systems[0] and len(model.humans) == 1
|
|
final = next(e for e in events if e.get("op") == "final")
|
|
assert "tool" in {n["type"] for n in final["graph"]["nodes"]}
|
|
|
|
|
|
async def test_generation_for_a_user_reports_what_it_could_not_fix(world, monkeypatch):
|
|
"""改不好的照实交付：自查失败、final 里列出 chat.not_allowed，不自动运行。"""
|
|
await _allow_chat(world)
|
|
model = _Scripted(world["source"])
|
|
model._lines = lambda repair: _Scripted(world["source"])._lines(False) if not repair else [
|
|
json.dumps({"op": "done", "explanation": "没改"}, ensure_ascii=False)]
|
|
events = await _generate(world["uma"], monkeypatch, model)
|
|
check = [e for e in events if e.get("op") == "check"][-1]
|
|
final = next(e for e in events if e.get("op") == "final")
|
|
assert check["status"] == "failed" and any(i["code"] == "chat.not_allowed" for i in check["issues"])
|
|
assert final["autorun"] is False and any(i["code"] == "chat.not_allowed" for i in final["issues"])
|
|
|
|
|
|
# ---------------------------------------------------------------- WebSocket
|
|
|
|
|
|
def _cookie(client) -> str:
|
|
return f"agentlab_session={client.cookies.get('agentlab_session')}"
|
|
|
|
|
|
async def _ws(path: str, cookie: str | None) -> list[dict]:
|
|
"""把一次 WebSocket 握手原样交给整个应用（中间件 + 路由级依赖 + 路由），记下发出去的帧。"""
|
|
from app.main import app
|
|
|
|
inbox: asyncio.Queue = asyncio.Queue()
|
|
await inbox.put({"type": "websocket.connect"})
|
|
sent: list[dict] = []
|
|
|
|
async def receive():
|
|
return await inbox.get()
|
|
|
|
async def send(message):
|
|
sent.append(message)
|
|
if message["type"] == "websocket.close":
|
|
await inbox.put({"type": "websocket.disconnect", "code": 1000})
|
|
|
|
headers = [(b"host", b"test")] + ([(b"cookie", cookie.encode())] if cookie else [])
|
|
scope = {"type": "websocket", "asgi": {"version": "3.0"}, "scheme": "ws", "http_version": "1.1", "path": path,
|
|
"raw_path": path.encode(), "query_string": b"", "root_path": "", "headers": headers,
|
|
"client": ("127.0.0.1", 50000), "server": ("test", 80), "subprotocols": [], "state": {}}
|
|
await asyncio.wait_for(app(scope, receive, send), 10)
|
|
return sent
|
|
|
|
|
|
async def _settled_run(**fields) -> str:
|
|
async with SessionLocal() as s:
|
|
run = Run(workflow_name="问数据", graph={}, input={}, status="succeeded", run_class="exploratory", **fields)
|
|
s.add(run)
|
|
await s.flush()
|
|
s.add(RunEvent(run_id=run.id, seq=1, type="run.started", ts=1.0, data={}))
|
|
await s.commit()
|
|
return run.id
|
|
|
|
|
|
def _refused(frames: list[dict]) -> bool:
|
|
return [f["type"] for f in frames] == ["websocket.close"] and frames[0].get("code") == 1008
|
|
|
|
|
|
async def test_websocket_is_judged_per_run_at_the_handshake(world):
|
|
mine = await _settled_run(started_by="uma", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
|
|
others = await _settled_run(started_by="bob", origin_kind="chat", options={MASK_RESULTS_OPTION: True})
|
|
raw = await _settled_run(started_by="uma", origin_kind="chat") # 关登录时自报同名留下的，没在边界遮过
|
|
uma, bob = _cookie(world["uma"]), _cookie(world["bob"])
|
|
|
|
assert _refused(await _ws(f"/api/runs/{mine}/stream", None)), "没登录的握手前关掉"
|
|
assert _refused(await _ws(f"/api/runs/{mine}/stream", uma)), "没开问数据时使用者连不了运行"
|
|
await _allow_chat(world)
|
|
frames = await _ws(f"/api/runs/{mine}/stream", uma)
|
|
assert frames[0]["type"] == "websocket.accept" and '"stream.end"' in frames[-2]["text"]
|
|
assert _refused(await _ws(f"/api/runs/{others}/stream", uma)), "别人的运行"
|
|
assert _refused(await _ws(f"/api/runs/{raw}/stream", uma)), "查询结果没在工具边界遮过的运行"
|
|
assert _refused(await _ws("/api/runs/nope/stream", uma))
|
|
assert (await _ws(f"/api/runs/{raw}/stream", bob))[0]["type"] == "websocket.accept", "搭建者照旧"
|
|
|
|
|
|
# ---------------------------------------------------------------- 数据源名单
|
|
|
|
|
|
async def test_source_names_for_the_chat_page(world):
|
|
uma, bob = world["uma"], world["bob"]
|
|
assert (await uma.get("/api/datasources/names")).status_code == 403, "没开问数据时不给使用者"
|
|
assert (await uma.get("/api/datasources")).status_code == 403
|
|
await _allow_chat(world)
|
|
body = (await uma.get("/api/datasources/names")).json()
|
|
assert [i["name"] for i in body["items"]] == [world["source"]], "只列启用着的"
|
|
assert set(body["items"][0]) == {"id", "name", "description", "kind"}, "连接信息不在名单里"
|
|
assert body["items"][0]["description"] == "客户" and body["items"][0]["kind"] == "sqlite"
|
|
assert (await uma.get("/api/datasources")).status_code == 403, "完整的数据源列表照旧只给搭建者"
|
|
assert (await bob.get("/api/datasources/names")).status_code == 200
|
|
|
|
|
|
async def test_source_names_with_login_off():
|
|
await auth_env.reset()
|
|
async with auth_env.client() as c:
|
|
r = await c.get("/api/datasources/names")
|
|
assert r.status_code == 200 and "items" in r.json()
|