|
1 | 1 | import asyncio |
| 2 | +from types import SimpleNamespace |
2 | 3 | from unittest.mock import AsyncMock, MagicMock |
3 | 4 |
|
4 | 5 | import pytest |
5 | 6 | import pytest_asyncio |
6 | 7 | from quart import Quart |
7 | 8 |
|
| 9 | +import astrbot.dashboard.routes.knowledge_base as knowledge_base_route_module |
8 | 10 | from astrbot.core import LogBroker |
9 | 11 | from astrbot.core.core_lifecycle import AstrBotCoreLifecycle |
10 | 12 | from astrbot.core.db.sqlite import SQLiteDatabase |
@@ -117,6 +119,55 @@ async def authenticated_header(app: Quart, core_lifecycle_td: AstrBotCoreLifecyc |
117 | 119 | return {"Authorization": f"Bearer {token}"} |
118 | 120 |
|
119 | 121 |
|
| 122 | +@pytest.mark.asyncio |
| 123 | +async def test_remove_deleted_kb_from_session_configs(monkeypatch: pytest.MonkeyPatch): |
| 124 | + updates = [] |
| 125 | + |
| 126 | + class FakeSharedPreferences: |
| 127 | + async def session_get(self, umo, key): |
| 128 | + assert umo is None |
| 129 | + assert key == "kb_config" |
| 130 | + return [ |
| 131 | + SimpleNamespace( |
| 132 | + scope_id="platform:GroupMessage:group!alice", |
| 133 | + value={"val": {"kb_ids": ["kb-old", "kb-keep"], "top_k": 3}}, |
| 134 | + ), |
| 135 | + SimpleNamespace( |
| 136 | + scope_id="platform:GroupMessage:group!bob", |
| 137 | + value={"val": {"kb_ids": ["kb-old"]}}, |
| 138 | + ), |
| 139 | + SimpleNamespace( |
| 140 | + scope_id="platform:FriendMessage:charlie", |
| 141 | + value={"val": {"kb_ids": ["kb-keep"]}}, |
| 142 | + ), |
| 143 | + SimpleNamespace( |
| 144 | + scope_id="platform:FriendMessage:broken", |
| 145 | + value={"val": {"kb_ids": "kb-old"}}, |
| 146 | + ), |
| 147 | + ] |
| 148 | + |
| 149 | + async def session_put(self, umo, key, value): |
| 150 | + updates.append((umo, key, value)) |
| 151 | + |
| 152 | + monkeypatch.setattr(knowledge_base_route_module, "sp", FakeSharedPreferences()) |
| 153 | + |
| 154 | + updated = await KnowledgeBaseRoute._remove_kb_from_session_configs("kb-old") |
| 155 | + |
| 156 | + assert updated == 2 |
| 157 | + assert updates == [ |
| 158 | + ( |
| 159 | + "platform:GroupMessage:group!alice", |
| 160 | + "kb_config", |
| 161 | + {"kb_ids": ["kb-keep"], "top_k": 3}, |
| 162 | + ), |
| 163 | + ( |
| 164 | + "platform:GroupMessage:group!bob", |
| 165 | + "kb_config", |
| 166 | + {"kb_ids": []}, |
| 167 | + ), |
| 168 | + ] |
| 169 | + |
| 170 | + |
120 | 171 | @pytest.mark.asyncio |
121 | 172 | async def test_import_documents( |
122 | 173 | app: Quart, authenticated_header: dict, core_lifecycle_td: AstrBotCoreLifecycle |
|
0 commit comments