test: 补充 3 个高优先级测试缺口 (52 tests)
U2: tests/unit/llm/test_migration.py (6 tests)
覆盖 migrate_api_keys_to_secrets 的 happy path/idempotent/empty/
source missing/write failure/dual source 场景。mock SecretsStore 隔离 I/O。
U3: tests/unit/test_memory_models.py (31 tests)
覆盖 EpisodeModel + ExperienceModel 的列类型/默认值/边界值。
源文件是 SQLAlchemy ORM 模型（非 Pydantic），按 ORM 等价测试。
跳过需 DB 连接的函数。
U4: tests/unit/test_chat_ws_routes.py (15 tests)
覆盖 WebSocket /api/v1/chat/ws/{session_id} 的连接/认证/消息
(ping/cancel/message/confirmation)/隔离/清理。mock _handle_chat_message。
所有测试通过，ruff 干净。预先存在的 OpenAICompatibleProvider 构造签名
和硬编码路径失败与本次改动无关。
This commit is contained in:
Chiguyong
parent
5f5c09835f
commit
e0a22e0447
|
|
@ -0,0 +1,223 @@
|
|||
"""U15 — migrate_api_keys_to_secrets 顶层函数单元测试。
|
||||
|
||||
测试 ``src/agentkit/llm/migration.py`` 中的 ``migrate_api_keys_to_secrets(config_path)``：
|
||||
读取 ``agentkit.yaml``，把每个 LLM provider 的 plaintext ``api_key`` 迁移到
|
||||
``SecretsStore`` 加密存储，并把更新后的 providers 段写回 YAML。
|
||||
|
||||
覆盖场景：
|
||||
1. Happy path — plaintext → secrets_store，YAML 写回，report 标记 migrated
|
||||
2. Idempotent — 已迁移配置再次调用返回 skipped，不重复加密
|
||||
3. Empty config — 无 providers 时返回 {}，YAML 仍可读
|
||||
4. Source missing — config_path 不存在时抛 FileNotFoundError（函数未捕获）
|
||||
5. Secrets store write failure — set_secret 抛异常时 status=failed，plaintext 保留（事务性）
|
||||
6. Dual source re-migrate — api_key_source=dual 时重新迁移覆盖旧 encrypted 值
|
||||
|
||||
隔离策略：
|
||||
- 用 ``tmp_path`` 隔离 YAML 文件 I/O
|
||||
- patch ``agentkit.channels.secrets.SecretsStore`` 注入可控实例（避免 env 依赖）
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from agentkit.channels.secrets import SecretsStore
|
||||
from agentkit.llm.migration import migrate_api_keys_to_secrets
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 测试辅助
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def _write_config(path: Path, providers: dict[str, dict[str, object]]) -> None:
|
||||
"""写一个最小的 agentkit.yaml（仅 llm.providers 段）。"""
|
||||
data: dict[str, dict[str, object]] = {"llm": {"providers": providers}}
|
||||
path.write_text(
|
||||
yaml.safe_dump(data, allow_unicode=True, sort_keys=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def _make_real_store() -> SecretsStore:
|
||||
"""构造一个真实可用的 SecretsStore（注入固定 master_key，不依赖 env）。"""
|
||||
return SecretsStore(master_key=b"x" * 32)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 1. Happy path
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_happy_path_migrates_plaintext_to_secrets_store(tmp_path: Path) -> None:
|
||||
"""plaintext api_key → secrets_store；YAML 写回，report 标记 migrated。"""
|
||||
config_path = tmp_path / "agentkit.yaml"
|
||||
_write_config(
|
||||
config_path,
|
||||
{
|
||||
"openai": {
|
||||
"api_key": "sk-secret",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"type": "openai",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||||
report = migrate_api_keys_to_secrets(config_path)
|
||||
|
||||
assert report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||||
|
||||
# YAML 应已更新：plaintext 清空，encrypted 写入，source 标记
|
||||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
pconf = raw["llm"]["providers"]["openai"]
|
||||
assert pconf["api_key"] == ""
|
||||
assert pconf["api_key_encrypted"] is not None
|
||||
assert pconf["api_key_source"] == "secrets_store"
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 2. Idempotent
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_idempotent_second_call_returns_skipped(tmp_path: Path) -> None:
|
||||
"""已迁移的配置再次调用 → skipped，不重复加密。"""
|
||||
config_path = tmp_path / "agentkit.yaml"
|
||||
_write_config(
|
||||
config_path,
|
||||
{
|
||||
"openai": {
|
||||
"api_key": "sk-secret",
|
||||
"base_url": "https://api.openai.com/v1",
|
||||
"type": "openai",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||||
first_report = migrate_api_keys_to_secrets(config_path)
|
||||
second_report = migrate_api_keys_to_secrets(config_path)
|
||||
|
||||
assert first_report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||||
assert second_report == {"openai": {"status": "skipped", "source": "secrets_store"}}
|
||||
|
||||
# 第二次调用后 YAML 仍一致（plaintext 仍为空）
|
||||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
pconf = raw["llm"]["providers"]["openai"]
|
||||
assert pconf["api_key"] == ""
|
||||
assert pconf["api_key_source"] == "secrets_store"
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 3. Empty config
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_empty_config_returns_empty_report(tmp_path: Path) -> None:
|
||||
"""无 providers 时返回 {}，YAML 仍可读且 providers 段为空。"""
|
||||
config_path = tmp_path / "agentkit.yaml"
|
||||
config_path.write_text("llm:\n providers: {}\n", encoding="utf-8")
|
||||
|
||||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||||
report = migrate_api_keys_to_secrets(config_path)
|
||||
|
||||
assert report == {}
|
||||
|
||||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
assert raw["llm"]["providers"] == {}
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 4. Source missing
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_source_missing_raises_file_not_found(tmp_path: Path) -> None:
|
||||
"""config_path 不存在时抛 FileNotFoundError（函数未捕获，向上传播）。"""
|
||||
missing = tmp_path / "nonexistent.yaml"
|
||||
with pytest.raises(FileNotFoundError):
|
||||
migrate_api_keys_to_secrets(missing)
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 5. Secrets store write failure
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_secrets_store_write_failure_marks_failed_and_retains_plaintext(
|
||||
tmp_path: Path,
|
||||
) -> None:
|
||||
"""set_secret 抛异常 → status=failed，plaintext 保留（事务性，不部分迁移）。"""
|
||||
config_path = tmp_path / "agentkit.yaml"
|
||||
_write_config(
|
||||
config_path,
|
||||
{
|
||||
"openai": {
|
||||
"api_key": "sk-secret",
|
||||
"base_url": "",
|
||||
"type": "openai",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
# mock store：set_secret 抛异常（get_secret 不会被调用到）
|
||||
fake_store = MagicMock(spec=SecretsStore)
|
||||
fake_store.set_secret = AsyncMock(side_effect=RuntimeError("redis down"))
|
||||
|
||||
with patch("agentkit.channels.secrets.SecretsStore", return_value=fake_store):
|
||||
report = migrate_api_keys_to_secrets(config_path)
|
||||
|
||||
assert report == {"openai": {"status": "failed", "source": "plaintext", "error": "redis down"}}
|
||||
|
||||
# plaintext 应保留（未清空），source 仍为 plaintext
|
||||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
pconf = raw["llm"]["providers"]["openai"]
|
||||
assert pconf["api_key"] == "sk-secret"
|
||||
assert pconf["api_key_source"] == "plaintext"
|
||||
|
||||
# set_secret 被调用一次（验证读 get_secret 未执行）
|
||||
fake_store.set_secret.assert_awaited_once()
|
||||
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 6. Dual source re-migrate (secret overwrite path)
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_dual_source_re_migrates_and_overwrites_encrypted(tmp_path: Path) -> None:
|
||||
"""api_key_source=dual（已有 encrypted 但 plaintext 保留）→ 重新迁移覆盖旧值。"""
|
||||
stale_encrypted = (
|
||||
'{"key":"llm:provider:openai:api_key","value":"old","nonce":"n","salt":"s",'
|
||||
'"key_id":"default","created_at":"","updated_at":""}'
|
||||
)
|
||||
config_path = tmp_path / "agentkit.yaml"
|
||||
_write_config(
|
||||
config_path,
|
||||
{
|
||||
"openai": {
|
||||
"api_key": "sk-secret",
|
||||
"base_url": "",
|
||||
"type": "openai",
|
||||
"api_key_encrypted": stale_encrypted,
|
||||
"api_key_source": "dual",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
with patch("agentkit.channels.secrets.SecretsStore", return_value=_make_real_store()):
|
||||
report = migrate_api_keys_to_secrets(config_path)
|
||||
|
||||
assert report == {"openai": {"status": "migrated", "source": "secrets_store"}}
|
||||
|
||||
raw = yaml.safe_load(config_path.read_text(encoding="utf-8"))
|
||||
pconf = raw["llm"]["providers"]["openai"]
|
||||
assert pconf["api_key"] == ""
|
||||
assert pconf["api_key_source"] == "secrets_store"
|
||||
# encrypted 应被新值覆盖（不等于原来的 stale 占位）
|
||||
assert pconf["api_key_encrypted"] != stale_encrypted
|
||||
assert pconf["api_key_encrypted"] is not None
|
||||
|
|
@ -0,0 +1,348 @@
|
|||
"""Unit tests for chat.py WebSocket route — message handling flows.
|
||||
|
||||
Tests the WebSocket endpoint at /api/v1/chat/ws/{session_id}, focusing on
|
||||
the message handling protocol: connected event, ping/pong, cancel,
|
||||
API key auth, session validation, happy-path message → final_answer,
|
||||
confirmation request/reply flow, invalid JSON handling, session isolation,
|
||||
and connection cleanup.
|
||||
|
||||
The LLM gateway and agent execution are mocked — these tests focus on the
|
||||
WebSocket protocol layer, not real LLM calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, WebSocket
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from agentkit.core.protocol import CancellationToken
|
||||
from agentkit.session.manager import SessionManager
|
||||
from agentkit.session.store import InMemorySessionStore
|
||||
|
||||
|
||||
# ── Stub handlers (match _handle_chat_message signature) ────────────────
|
||||
#
|
||||
# 10 params — all required by the call site in chat_websocket which passes
|
||||
# them positionally. Full signature avoids *args (no Any usage).
|
||||
|
||||
|
||||
async def _stub_final_answer(
|
||||
websocket: WebSocket,
|
||||
session_id: str,
|
||||
content: str,
|
||||
sm: SessionManager,
|
||||
cancellation_token: CancellationToken,
|
||||
pending_replies: dict[str, asyncio.Future],
|
||||
pending_confirmations: dict[str, asyncio.Future] | None = None,
|
||||
pending_spec_reviews: dict[str, asyncio.Future] | None = None,
|
||||
pending_autonomy_resumes: dict[str, asyncio.Future] | None = None,
|
||||
model_override: str | None = None,
|
||||
) -> None:
|
||||
"""Stub that immediately sends a final_answer event."""
|
||||
await websocket.send_json(
|
||||
{"type": "final_answer", "content": f"echo: {content}", "is_final": True}
|
||||
)
|
||||
|
||||
|
||||
async def _stub_with_confirmation(
|
||||
websocket: WebSocket,
|
||||
session_id: str,
|
||||
content: str,
|
||||
sm: SessionManager,
|
||||
cancellation_token: CancellationToken,
|
||||
pending_replies: dict[str, asyncio.Future],
|
||||
pending_confirmations: dict[str, asyncio.Future] | None = None,
|
||||
pending_spec_reviews: dict[str, asyncio.Future] | None = None,
|
||||
pending_autonomy_resumes: dict[str, asyncio.Future] | None = None,
|
||||
model_override: str | None = None,
|
||||
) -> None:
|
||||
"""Stub: send confirmation_request, wait for reply, send final_answer."""
|
||||
conf_id = "test-conf-1"
|
||||
# Register future BEFORE sending request so the WS loop can resolve it
|
||||
# even if the reply arrives before we await (race-safe).
|
||||
loop = asyncio.get_running_loop()
|
||||
future: asyncio.Future[bool] = loop.create_future()
|
||||
if pending_confirmations is not None:
|
||||
pending_confirmations[conf_id] = future
|
||||
await websocket.send_json(
|
||||
{
|
||||
"type": "confirmation_request",
|
||||
"data": {
|
||||
"confirmation_id": conf_id,
|
||||
"command": "rm -rf /tmp/test",
|
||||
"reason": "dangerous command",
|
||||
},
|
||||
}
|
||||
)
|
||||
try:
|
||||
approved = await asyncio.wait_for(future, timeout=10.0)
|
||||
except asyncio.TimeoutError:
|
||||
approved = False
|
||||
await websocket.send_json(
|
||||
{"type": "final_answer", "content": f"approved={approved}", "is_final": True}
|
||||
)
|
||||
|
||||
|
||||
# ── Fixtures ────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_with_chat_ws() -> FastAPI:
|
||||
"""Create a FastAPI app with Chat routes and mocked dependencies."""
|
||||
from agentkit.server.routes.chat import router
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(router, prefix="/api/v1")
|
||||
app.state.session_manager = SessionManager(store=InMemorySessionStore())
|
||||
app.state.llm_gateway = MagicMock()
|
||||
app.state.agent_pool = MagicMock()
|
||||
app.state.server_config = MagicMock()
|
||||
app.state.server_config.api_key = None
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client(app_with_chat_ws: FastAPI) -> TestClient:
|
||||
return TestClient(app_with_chat_ws)
|
||||
|
||||
|
||||
def _create_session(client: TestClient, agent_name: str = "test-agent") -> str:
|
||||
"""Create a chat session via REST and return its id."""
|
||||
resp = client.post("/api/v1/chat/sessions", json={"agent_name": agent_name})
|
||||
assert resp.status_code == 200
|
||||
return resp.json()["session_id"]
|
||||
|
||||
|
||||
# ── Connection & session validation ─────────────────────────────────────
|
||||
|
||||
|
||||
class TestWSConnection:
|
||||
"""Connection establishment and session validation."""
|
||||
|
||||
def test_connected_event_with_session_id(self, client: TestClient) -> None:
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "connected"
|
||||
assert msg["session_id"] == session_id
|
||||
|
||||
def test_session_not_found_rejected(self, client: TestClient) -> None:
|
||||
with client.websocket_connect("/api/v1/chat/ws/nonexistent-session") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "error"
|
||||
assert "not found" in msg["data"]["message"].lower()
|
||||
|
||||
def test_closed_session_rejected(self, client: TestClient) -> None:
|
||||
session_id = _create_session(client)
|
||||
client.delete(f"/api/v1/chat/sessions/{session_id}")
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "error"
|
||||
assert "closed" in msg["data"]["message"].lower()
|
||||
|
||||
|
||||
# ── Authentication ──────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestWSAuth:
|
||||
"""API key authentication via ?api_key= query param."""
|
||||
|
||||
def test_missing_api_key_rejected(self, app_with_chat_ws: FastAPI) -> None:
|
||||
app_with_chat_ws.state.server_config.api_key = "secret123"
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "error"
|
||||
assert "api_key" in msg["data"]["message"].lower()
|
||||
|
||||
def test_wrong_api_key_rejected(self, app_with_chat_ws: FastAPI) -> None:
|
||||
app_with_chat_ws.state.server_config.api_key = "secret123"
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}?api_key=wrong") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "error"
|
||||
assert "api_key" in msg["data"]["message"].lower()
|
||||
|
||||
def test_valid_api_key_accepted(self, app_with_chat_ws: FastAPI) -> None:
|
||||
app_with_chat_ws.state.server_config.api_key = "secret123"
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}?api_key=secret123") as ws:
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "connected"
|
||||
assert msg["session_id"] == session_id
|
||||
|
||||
|
||||
# ── Message handling ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestWSMessages:
|
||||
"""Message type handling: ping, cancel, message, confirmation, invalid JSON."""
|
||||
|
||||
def test_ping_returns_pong(self, client: TestClient) -> None:
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "ping"})
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "pong"
|
||||
|
||||
def test_cancel_sends_result_and_stays_open(self, client: TestClient) -> None:
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "cancel"})
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "result"
|
||||
assert msg["data"]["status"] == "cancelled"
|
||||
# Connection stays open — ping still works after cancel
|
||||
ws.send_json({"type": "ping"})
|
||||
assert ws.receive_json()["type"] == "pong"
|
||||
|
||||
def test_message_happy_path(
|
||||
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from agentkit.server.routes import chat as chat_module
|
||||
|
||||
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_final_answer)
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "message", "content": "hello"})
|
||||
final = ws.receive_json()
|
||||
assert final["type"] == "final_answer"
|
||||
assert final["content"] == "echo: hello"
|
||||
assert final["is_final"] is True
|
||||
|
||||
def test_confirmation_flow(
|
||||
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Server sends confirmation_request → client replies → server continues."""
|
||||
from agentkit.server.routes import chat as chat_module
|
||||
|
||||
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_with_confirmation)
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "message", "content": "do something risky"})
|
||||
|
||||
conf_req = ws.receive_json()
|
||||
assert conf_req["type"] == "confirmation_request"
|
||||
assert conf_req["data"]["confirmation_id"] == "test-conf-1"
|
||||
|
||||
ws.send_json(
|
||||
{
|
||||
"type": "confirmation_reply",
|
||||
"confirmation_id": "test-conf-1",
|
||||
"approved": True,
|
||||
}
|
||||
)
|
||||
|
||||
final = ws.receive_json()
|
||||
assert final["type"] == "final_answer"
|
||||
assert final["content"] == "approved=True"
|
||||
|
||||
def test_confirmation_flow_rejected(
|
||||
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Client rejects confirmation → final_answer reflects approved=False."""
|
||||
from agentkit.server.routes import chat as chat_module
|
||||
|
||||
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_with_confirmation)
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "message", "content": "risky"})
|
||||
|
||||
conf_req = ws.receive_json()
|
||||
assert conf_req["type"] == "confirmation_request"
|
||||
|
||||
ws.send_json(
|
||||
{
|
||||
"type": "confirmation_reply",
|
||||
"confirmation_id": "test-conf-1",
|
||||
"approved": False,
|
||||
}
|
||||
)
|
||||
|
||||
final = ws.receive_json()
|
||||
assert final["type"] == "final_answer"
|
||||
assert final["content"] == "approved=False"
|
||||
|
||||
def test_invalid_json_ignored(self, client: TestClient) -> None:
|
||||
"""Invalid JSON text is silently ignored — connection stays open.
|
||||
|
||||
Note: the implementation catches JSONDecodeError with `continue`
|
||||
rather than sending an error event. This is the actual behavior.
|
||||
"""
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_text("not valid json {{{")
|
||||
# Connection stays open — ping still works
|
||||
ws.send_json({"type": "ping"})
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "pong"
|
||||
|
||||
def test_unknown_message_type_ignored(self, client: TestClient) -> None:
|
||||
"""Unknown message type is silently ignored — no error sent."""
|
||||
session_id = _create_session(client)
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
ws.send_json({"type": "unknown_type", "data": "whatever"})
|
||||
# Connection stays open — ping still works
|
||||
ws.send_json({"type": "ping"})
|
||||
msg = ws.receive_json()
|
||||
assert msg["type"] == "pong"
|
||||
|
||||
|
||||
# ── Isolation & cleanup ────────────────────────────────────────────────
|
||||
|
||||
|
||||
class TestWSIsolationAndCleanup:
|
||||
"""Session isolation and connection cleanup."""
|
||||
|
||||
def test_session_isolation(self, client: TestClient) -> None:
|
||||
"""Different session_ids get separate connected events."""
|
||||
sid1 = _create_session(client)
|
||||
sid2 = _create_session(client)
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{sid1}") as ws1:
|
||||
msg1 = ws1.receive_json()
|
||||
assert msg1["session_id"] == sid1
|
||||
ws1.send_json({"type": "ping"})
|
||||
assert ws1.receive_json()["type"] == "pong"
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{sid2}") as ws2:
|
||||
msg2 = ws2.receive_json()
|
||||
assert msg2["session_id"] == sid2
|
||||
ws2.send_json({"type": "ping"})
|
||||
assert ws2.receive_json()["type"] == "pong"
|
||||
|
||||
def test_connection_cleanup_on_disconnect(self, app_with_chat_ws: FastAPI) -> None:
|
||||
"""After disconnect, chat_manager no longer tracks the connection."""
|
||||
from agentkit.server.routes.chat import chat_manager
|
||||
|
||||
client = TestClient(app_with_chat_ws)
|
||||
session_id = _create_session(client)
|
||||
|
||||
assert len(chat_manager.get_connections(session_id)) == 0
|
||||
|
||||
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
||||
ws.receive_json() # connected
|
||||
assert len(chat_manager.get_connections(session_id)) == 1
|
||||
|
||||
# After disconnect, the finally block removes the connection
|
||||
assert len(chat_manager.get_connections(session_id)) == 0
|
||||
|
|
@ -0,0 +1,282 @@
|
|||
"""Unit tests for SQLAlchemy ORM models in agentkit.memory.models.
|
||||
|
||||
Covers the two public ORM model classes — EpisodeModel and ExperienceModel:
|
||||
table mapping, column types/constraints, Python-side default callables,
|
||||
explicit-value construction, and edge cases (empty/long/unicode values).
|
||||
|
||||
Note: memory/models.py contains SQLAlchemy ORM models (not Pydantic), so
|
||||
Pydantic-specific scenarios (ValidationError / model_dump / model_validate)
|
||||
are replaced by ORM equivalents (column introspection + default callable
|
||||
invocation + attribute assignment).
|
||||
|
||||
DB-dependent helpers (create_episodic_session_factory,
|
||||
create_experience_session_factory, ensure_episodic_table) are NOT tested
|
||||
here — they require a live PostgreSQL + asyncpg connection.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from sqlalchemy import DateTime, Float, String, Text
|
||||
from sqlalchemy.dialects.postgresql import JSONB
|
||||
|
||||
from agentkit.memory.models import EpisodeModel, ExperienceModel
|
||||
|
||||
|
||||
class TestEpisodeModel:
|
||||
def test_tablename(self):
|
||||
assert EpisodeModel.__tablename__ == "episodic_memories"
|
||||
|
||||
def test_column_types(self):
|
||||
cols = EpisodeModel.__table__.c
|
||||
assert isinstance(cols.id.type, String)
|
||||
assert isinstance(cols.agent_name.type, String)
|
||||
assert isinstance(cols.task_type.type, String)
|
||||
assert isinstance(cols.input_summary.type, Text)
|
||||
assert isinstance(cols.output_summary.type, Text)
|
||||
assert isinstance(cols.outcome.type, String)
|
||||
assert isinstance(cols.quality_score.type, Float)
|
||||
assert isinstance(cols.reflection.type, Text)
|
||||
assert isinstance(cols.embedding.type, Text)
|
||||
assert isinstance(cols.metadata.type, JSONB)
|
||||
assert isinstance(cols.created_at.type, DateTime)
|
||||
|
||||
def test_primary_key(self):
|
||||
assert EpisodeModel.__table__.c.id.primary_key is True
|
||||
|
||||
def test_indexed_columns(self):
|
||||
cols = EpisodeModel.__table__.c
|
||||
assert cols.agent_name.index is True
|
||||
assert cols.task_type.index is True
|
||||
assert cols.created_at.index is True
|
||||
|
||||
def test_nullable_columns(self):
|
||||
cols = EpisodeModel.__table__.c
|
||||
assert cols.embedding.nullable is True
|
||||
assert cols.metadata.nullable is True
|
||||
|
||||
def test_metadata_attribute_maps_to_metadata_column(self):
|
||||
# Python attribute is `metadata_` (avoids clash with Base.metadata);
|
||||
# the underlying DB column name is "metadata".
|
||||
assert hasattr(EpisodeModel, "metadata_")
|
||||
assert EpisodeModel.__table__.c.metadata.name == "metadata"
|
||||
|
||||
def test_scalar_defaults(self):
|
||||
cols = EpisodeModel.__table__.c
|
||||
assert cols.input_summary.default.arg == ""
|
||||
assert cols.output_summary.default.arg == ""
|
||||
assert cols.reflection.default.arg == ""
|
||||
assert cols.outcome.default.arg == "success"
|
||||
assert cols.quality_score.default.arg == 0.5
|
||||
|
||||
def test_callable_default_id_is_uuid_string(self):
|
||||
# SQLAlchemy wraps Python-side default callables to receive an execution
|
||||
# context; pass None for callables that don't inspect it.
|
||||
id_default = EpisodeModel.__table__.c.id.default.arg
|
||||
val = id_default(None)
|
||||
assert isinstance(val, str)
|
||||
uuid.UUID(val) # raises ValueError if not a valid uuid4 string
|
||||
|
||||
def test_callable_default_created_at_is_utc(self):
|
||||
ts_default = EpisodeModel.__table__.c.created_at.default.arg
|
||||
ts = ts_default(None)
|
||||
assert isinstance(ts, datetime)
|
||||
assert ts.tzinfo is not None
|
||||
assert ts.utcoffset() == timezone.utc.utcoffset(None)
|
||||
|
||||
def test_default_id_is_unique_across_calls(self):
|
||||
id_default = EpisodeModel.__table__.c.id.default.arg
|
||||
ids = {id_default(None) for _ in range(100)}
|
||||
assert len(ids) == 100
|
||||
|
||||
def test_construct_with_explicit_values(self):
|
||||
ts = datetime(2025, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
|
||||
m = EpisodeModel(
|
||||
id="ep-1",
|
||||
agent_name="agent-a",
|
||||
task_type="analysis",
|
||||
input_summary="in",
|
||||
output_summary="out",
|
||||
outcome="failure",
|
||||
quality_score=0.9,
|
||||
reflection="ref",
|
||||
embedding="[0.1, 0.2]",
|
||||
metadata_={"k": "v"},
|
||||
created_at=ts,
|
||||
)
|
||||
assert m.id == "ep-1"
|
||||
assert m.agent_name == "agent-a"
|
||||
assert m.task_type == "analysis"
|
||||
assert m.input_summary == "in"
|
||||
assert m.output_summary == "out"
|
||||
assert m.outcome == "failure"
|
||||
assert m.quality_score == 0.9
|
||||
assert m.reflection == "ref"
|
||||
assert m.embedding == "[0.1, 0.2]"
|
||||
assert m.metadata_ == {"k": "v"}
|
||||
assert m.created_at == ts
|
||||
|
||||
def test_construct_minimal_yields_none_attrs(self):
|
||||
# ORM Python-side defaults fire at flush (INSERT) time, not instantiation,
|
||||
# so unset attributes read back as None.
|
||||
m = EpisodeModel()
|
||||
assert m.id is None
|
||||
assert m.agent_name is None
|
||||
assert m.created_at is None
|
||||
assert m.metadata_ is None
|
||||
|
||||
def test_edge_empty_embedding_and_text(self):
|
||||
m = EpisodeModel(embedding="", reflection="", input_summary="")
|
||||
assert m.embedding == ""
|
||||
assert m.reflection == ""
|
||||
assert m.input_summary == ""
|
||||
|
||||
def test_edge_empty_vector_json(self):
|
||||
m = EpisodeModel(embedding="[]")
|
||||
assert m.embedding == "[]"
|
||||
|
||||
def test_edge_long_text(self):
|
||||
long_text = "x" * 10_000
|
||||
m = EpisodeModel(input_summary=long_text, output_summary=long_text, reflection=long_text)
|
||||
assert len(m.input_summary) == 10_000
|
||||
assert len(m.output_summary) == 10_000
|
||||
assert len(m.reflection) == 10_000
|
||||
|
||||
def test_edge_unicode_and_special_characters(self):
|
||||
m = EpisodeModel(
|
||||
input_summary="中文 + emoji 🚀 + quote \" ' \n\t",
|
||||
metadata_={"unicode": "你好", "null": None, "nested": {"a": [1, 2]}},
|
||||
)
|
||||
assert "🚀" in m.input_summary
|
||||
assert "中文" in m.input_summary
|
||||
assert m.metadata_["unicode"] == "你好"
|
||||
assert m.metadata_["null"] is None
|
||||
assert m.metadata_["nested"] == {"a": [1, 2]}
|
||||
|
||||
def test_edge_empty_metadata_dict(self):
|
||||
m = EpisodeModel(metadata_={})
|
||||
assert m.metadata_ == {}
|
||||
|
||||
|
||||
class TestExperienceModel:
|
||||
def test_tablename(self):
|
||||
assert ExperienceModel.__tablename__ == "task_experiences"
|
||||
|
||||
def test_column_types(self):
|
||||
cols = ExperienceModel.__table__.c
|
||||
assert isinstance(cols.id.type, String)
|
||||
assert isinstance(cols.task_type.type, String)
|
||||
assert isinstance(cols.goal.type, Text)
|
||||
assert isinstance(cols.steps_summary.type, Text)
|
||||
assert isinstance(cols.outcome.type, String)
|
||||
assert isinstance(cols.duration_seconds.type, Float)
|
||||
assert isinstance(cols.success_rate.type, Float)
|
||||
assert isinstance(cols.failure_reasons.type, JSONB)
|
||||
assert isinstance(cols.optimization_tips.type, JSONB)
|
||||
assert isinstance(cols.embedding.type, Text)
|
||||
assert isinstance(cols.created_at.type, DateTime)
|
||||
|
||||
def test_primary_key_and_indexes(self):
|
||||
cols = ExperienceModel.__table__.c
|
||||
assert cols.id.primary_key is True
|
||||
assert cols.task_type.index is True
|
||||
assert cols.created_at.index is True
|
||||
|
||||
def test_nullable_embedding(self):
|
||||
assert ExperienceModel.__table__.c.embedding.nullable is True
|
||||
|
||||
def test_scalar_defaults(self):
|
||||
cols = ExperienceModel.__table__.c
|
||||
assert cols.goal.default.arg == ""
|
||||
assert cols.steps_summary.default.arg == ""
|
||||
assert cols.outcome.default.arg == "success"
|
||||
assert cols.duration_seconds.default.arg == 0.0
|
||||
assert cols.success_rate.default.arg == 1.0
|
||||
|
||||
def test_callable_list_defaults(self):
|
||||
cols = ExperienceModel.__table__.c
|
||||
# default=list (the type itself, callable → empty list when invoked)
|
||||
assert cols.failure_reasons.default.is_callable
|
||||
assert cols.optimization_tips.default.is_callable
|
||||
assert cols.failure_reasons.default.arg(None) == []
|
||||
assert cols.optimization_tips.default.arg(None) == []
|
||||
|
||||
def test_callable_default_id_is_uuid_string(self):
|
||||
id_default = ExperienceModel.__table__.c.id.default.arg
|
||||
val = id_default(None)
|
||||
assert isinstance(val, str)
|
||||
uuid.UUID(val)
|
||||
|
||||
def test_callable_default_created_at_is_utc(self):
|
||||
ts_default = ExperienceModel.__table__.c.created_at.default.arg
|
||||
ts = ts_default(None)
|
||||
assert isinstance(ts, datetime)
|
||||
assert ts.tzinfo is not None
|
||||
assert ts.utcoffset() == timezone.utc.utcoffset(None)
|
||||
|
||||
def test_default_id_is_unique_across_calls(self):
|
||||
id_default = ExperienceModel.__table__.c.id.default.arg
|
||||
ids = {id_default(None) for _ in range(100)}
|
||||
assert len(ids) == 100
|
||||
|
||||
def test_construct_with_explicit_values(self):
|
||||
ts = datetime(2025, 6, 7, 8, 9, 10, tzinfo=timezone.utc)
|
||||
m = ExperienceModel(
|
||||
id="exp-1",
|
||||
task_type="build",
|
||||
goal="ship feature",
|
||||
steps_summary="step1; step2",
|
||||
outcome="partial",
|
||||
duration_seconds=42.5,
|
||||
success_rate=0.66,
|
||||
failure_reasons=["timeout", "oom"],
|
||||
optimization_tips=["retry", "cache"],
|
||||
embedding="[0.3, 0.4]",
|
||||
created_at=ts,
|
||||
)
|
||||
assert m.id == "exp-1"
|
||||
assert m.task_type == "build"
|
||||
assert m.goal == "ship feature"
|
||||
assert m.steps_summary == "step1; step2"
|
||||
assert m.outcome == "partial"
|
||||
assert m.duration_seconds == 42.5
|
||||
assert m.success_rate == 0.66
|
||||
assert m.failure_reasons == ["timeout", "oom"]
|
||||
assert m.optimization_tips == ["retry", "cache"]
|
||||
assert m.embedding == "[0.3, 0.4]"
|
||||
assert m.created_at == ts
|
||||
|
||||
def test_construct_minimal_yields_none_attrs(self):
|
||||
m = ExperienceModel()
|
||||
assert m.id is None
|
||||
assert m.task_type is None
|
||||
assert m.created_at is None
|
||||
assert m.failure_reasons is None
|
||||
assert m.optimization_tips is None
|
||||
|
||||
def test_edge_empty_lists_and_embedding(self):
|
||||
m = ExperienceModel(
|
||||
failure_reasons=[],
|
||||
optimization_tips=[],
|
||||
embedding="",
|
||||
)
|
||||
assert m.failure_reasons == []
|
||||
assert m.optimization_tips == []
|
||||
assert m.embedding == ""
|
||||
|
||||
def test_edge_long_text(self):
|
||||
long_text = "y" * 10_000
|
||||
m = ExperienceModel(goal=long_text, steps_summary=long_text)
|
||||
assert len(m.goal) == 10_000
|
||||
assert len(m.steps_summary) == 10_000
|
||||
|
||||
def test_edge_unicode_and_special_characters(self):
|
||||
m = ExperienceModel(
|
||||
goal="中文 🎯 ' \" \n\t",
|
||||
failure_reasons=["错误", "🚨"],
|
||||
optimization_tips=["优化"],
|
||||
)
|
||||
assert "🎯" in m.goal
|
||||
assert "中文" in m.goal
|
||||
assert m.failure_reasons == ["错误", "🚨"]
|
||||
assert m.optimization_tips == ["优化"]
|
||||
Loading…
Reference in New Issue