mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-10-31 04:45:23 +08:00
667 lines
28 KiB
Python
667 lines
28 KiB
Python
"""Real plugin/SDK API + synthetic Host KV tests (not a live LangBot DB).
|
|
Run from HumanTakeover: ../.venv/bin/python -m unittest discover -s tests -v
|
|
"""
|
|
|
|
import asyncio
|
|
import base64
|
|
import copy
|
|
import json
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from langbot_plugin.entities.io.actions.enums import PluginToRuntimeAction as Action
|
|
from main import HumanTakeover
|
|
|
|
|
|
class StorageHost:
|
|
def __init__(self):
|
|
self.data = {}
|
|
self.writes = []
|
|
self.reads = []
|
|
self.failure = None
|
|
self.write_started = None
|
|
self.release_write = None
|
|
|
|
async def call_action(self, action, data):
|
|
key = data.get("key")
|
|
if self.failure and self.failure(action, key):
|
|
raise OSError("fixture storage unavailable")
|
|
if action == Action.GET_PLUGIN_STORAGE_KEYS:
|
|
return {"keys": list(self.data)}
|
|
if action == Action.GET_PLUGIN_STORAGE:
|
|
self.reads.append(key)
|
|
return {"value_base64": base64.b64encode(self.data[key]).decode()}
|
|
if action == Action.SET_PLUGIN_STORAGE:
|
|
if len(json.dumps(data).encode()) + 1024 > 16 * 1024 * 1024:
|
|
raise ValueError("fixture 16 MiB encoded frame limit")
|
|
value = base64.b64decode(data["value_base64"])
|
|
self.writes.append((key, value))
|
|
if self.write_started is not None:
|
|
started, release = self.write_started, self.release_write
|
|
self.write_started = None
|
|
started.set()
|
|
await release.wait()
|
|
await asyncio.sleep(0)
|
|
self.data[key] = value
|
|
return {}
|
|
if action == Action.DELETE_PLUGIN_STORAGE:
|
|
del self.data[key]
|
|
return {}
|
|
raise AssertionError(action)
|
|
|
|
|
|
class DetachedHost(StorageHost):
|
|
"""The host commit outlives cancellation/loss of the RPC waiter."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pause_action = None
|
|
self.started = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
self.remote = None
|
|
self.error = None
|
|
|
|
async def call_action(self, action, data):
|
|
if action != self.pause_action:
|
|
return await super().call_action(action, data)
|
|
self.pause_action = None
|
|
|
|
async def commit():
|
|
self.started.set()
|
|
await self.release.wait()
|
|
return await super(DetachedHost, self).call_action(action, data)
|
|
|
|
self.remote = asyncio.create_task(commit())
|
|
await self.started.wait()
|
|
if self.error:
|
|
raise self.error("remote outcome unknown")
|
|
return await asyncio.shield(self.remote)
|
|
|
|
|
|
class PausableHost(StorageHost):
|
|
"""Holds one chosen action in flight, then fails it after the pause.
|
|
|
|
Unlike ``DetachedHost`` the caller stays blocked until the test releases it,
|
|
so another caller can be queued behind the installation lock first.
|
|
"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.pause_action = None
|
|
self.paused = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
self.pause_error = None
|
|
|
|
async def call_action(self, action, data):
|
|
if action != self.pause_action:
|
|
return await super().call_action(action, data)
|
|
self.pause_action = None
|
|
self.paused.set()
|
|
await self.release.wait()
|
|
if self.pause_error is not None:
|
|
raise self.pause_error("remote outcome unknown")
|
|
return await super().call_action(action, data)
|
|
|
|
|
|
async def open_plugin(host):
|
|
plugin = HumanTakeover()
|
|
plugin.config = {}
|
|
plugin.plugin_runtime_handler = host
|
|
await plugin.initialize()
|
|
# Shared placement loads tenant state on demand, per invocation.
|
|
await plugin.load_state()
|
|
return plugin
|
|
|
|
|
|
async def record(plugin, key="group_1", content="hello"):
|
|
await plugin.record_message(
|
|
session_key=key,
|
|
session_type="group",
|
|
target_id=key,
|
|
bot_uuid="fixture-bot",
|
|
session_name=key,
|
|
role="user",
|
|
sender_id="1",
|
|
sender_name="测试",
|
|
content_type="text",
|
|
content=content,
|
|
)
|
|
|
|
|
|
def seed_legacy(host, count=2, content="历史消息"):
|
|
sessions = {
|
|
f"group_{i}": {
|
|
"session_key": f"group_{i}",
|
|
"name": f"Group {i}",
|
|
"takeover": {"active": False},
|
|
"custom": {"keep": True},
|
|
}
|
|
for i in range(count)
|
|
}
|
|
messages = {
|
|
key: [{"role": "user", "content": content, "ts": 0}] for key in sessions
|
|
}
|
|
messages["person_orphan"] = [{"content": "do not discard"}]
|
|
host.data["ht_sessions"] = json.dumps(sessions, ensure_ascii=False).encode()
|
|
host.data["ht_messages"] = json.dumps(messages, ensure_ascii=False).encode()
|
|
return sessions, messages
|
|
|
|
|
|
class StorageTests(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self):
|
|
# Expected failure-injection logs are asserted via exceptions, not test stderr.
|
|
logger_patch = patch("main.logger")
|
|
logger_patch.start()
|
|
self.addCleanup(logger_patch.stop)
|
|
asyncio.get_running_loop().slow_callback_duration = 1
|
|
self.host = StorageHost()
|
|
self.plugin = await open_plugin(self.host)
|
|
self.host.writes.clear()
|
|
|
|
async def test_legacy_migration_preserves_exact_history_and_skips_legacy_on_restart(
|
|
self,
|
|
):
|
|
host = StorageHost()
|
|
sessions, messages = seed_legacy(host)
|
|
plugin = await open_plugin(host)
|
|
self.assertEqual((plugin.sessions, plugin.messages), (sessions, messages))
|
|
self.assertTrue(host.writes, "legacy data must migrate before normal writes")
|
|
host.reads.clear()
|
|
restarted = await open_plugin(host)
|
|
self.assertEqual((restarted.sessions, restarted.messages), (sessions, messages))
|
|
self.assertNotIn("ht_messages", host.reads)
|
|
self.assertNotIn("ht_sessions", host.reads)
|
|
|
|
async def test_update_writes_one_session_not_global_history(self):
|
|
await record(self.plugin, "group_1")
|
|
await record(self.plugin, "group_2", "unrelated")
|
|
before = dict(self.host.data)
|
|
self.host.writes.clear()
|
|
await record(self.plugin, "group_1", "new")
|
|
self.assertEqual(len(self.host.writes), 1)
|
|
key, value = self.host.writes[0]
|
|
self.assertNotIn(key, {"ht_messages", "ht_sessions"})
|
|
self.assertNotIn(b"unrelated", value)
|
|
self.assertEqual(sum(before.get(k) != v for k, v in self.host.data.items()), 1)
|
|
restarted = await open_plugin(self.host)
|
|
self.assertEqual(restarted.messages, self.plugin.messages)
|
|
self.assertEqual(restarted.sessions, self.plugin.sessions)
|
|
|
|
async def test_aggregate_history_above_frame_limit_can_record(self):
|
|
for i in range(1200):
|
|
key = f"group_{i}"
|
|
self.plugin._ensure_session(key, "group", key, "fixture-bot", key)
|
|
self.plugin.messages[key] = [{"content": "x" * 1024} for _ in range(11)]
|
|
await record(self.plugin, "group_0", "persist-me")
|
|
restarted = await open_plugin(self.host)
|
|
self.assertIn(
|
|
"group_0", restarted.messages, "successful record must survive restart"
|
|
)
|
|
self.assertEqual(restarted.messages["group_0"], self.plugin.messages["group_0"])
|
|
self.assertEqual(len(self.host.writes), 1)
|
|
self.assertLess(len(self.host.writes[0][1]), 32 * 1024)
|
|
|
|
async def test_concurrent_same_session_updates_cannot_commit_out_of_order(self):
|
|
await record(self.plugin)
|
|
self.host.write_started = started = asyncio.Event()
|
|
self.host.release_write = release = asyncio.Event()
|
|
first = asyncio.create_task(record(self.plugin, content="first"))
|
|
await started.wait()
|
|
second = asyncio.create_task(record(self.plugin, content="second"))
|
|
await asyncio.sleep(0.01)
|
|
release.set()
|
|
await asyncio.gather(first, second)
|
|
restarted = await open_plugin(self.host)
|
|
self.assertEqual(restarted.messages, self.plugin.messages)
|
|
self.assertEqual(restarted.sessions, self.plugin.sessions)
|
|
|
|
async def test_concurrent_sessions_and_metadata_read_back(self):
|
|
await asyncio.gather(*(record(self.plugin, f"group_{i}") for i in range(10)))
|
|
await asyncio.gather(
|
|
self.plugin.set_takeover("group_1", True),
|
|
self.plugin.mark_unread("group_1", "help"),
|
|
self.plugin.touch_human_response("group_2"),
|
|
)
|
|
restarted = await open_plugin(self.host)
|
|
self.assertEqual(restarted.sessions, self.plugin.sessions)
|
|
self.assertEqual(restarted.messages, self.plugin.messages)
|
|
|
|
async def test_write_failure_raises_and_rolls_back_memory(self):
|
|
await record(self.plugin)
|
|
before = copy.deepcopy((self.plugin.sessions, self.plugin.messages))
|
|
self.host.failure = lambda action, key: action == Action.SET_PLUGIN_STORAGE
|
|
with self.assertRaisesRegex(OSError, "storage unavailable"):
|
|
await record(self.plugin, content="not saved")
|
|
self.assertEqual((self.plugin.sessions, self.plugin.messages), before)
|
|
self.host.failure = None
|
|
# This fixture proves no request remains in flight; production must reconcile.
|
|
await self.plugin.reconcile()
|
|
await record(self.plugin, content="retry")
|
|
restarted = await open_plugin(self.host)
|
|
self.assertEqual(restarted.messages, self.plugin.messages)
|
|
|
|
async def test_failed_new_session_is_not_left_in_memory(self):
|
|
self.host.failure = lambda action, key: action == Action.SET_PLUGIN_STORAGE
|
|
with self.assertRaises(OSError):
|
|
await record(self.plugin)
|
|
self.assertEqual(self.plugin.sessions, {})
|
|
self.assertEqual(self.plugin.messages, {})
|
|
|
|
async def test_metadata_failure_does_not_report_success(self):
|
|
await record(self.plugin)
|
|
before = copy.deepcopy(self.plugin.sessions)
|
|
self.host.failure = lambda action, key: action == Action.SET_PLUGIN_STORAGE
|
|
with self.assertRaises(OSError):
|
|
await self.plugin.set_takeover("group_1", True)
|
|
self.assertEqual(self.plugin.sessions, before)
|
|
|
|
async def test_key_listing_failure_prevents_initialization(self):
|
|
self.host.failure = lambda action, key: action == Action.GET_PLUGIN_STORAGE_KEYS
|
|
with self.assertRaises(OSError):
|
|
await open_plugin(self.host)
|
|
|
|
async def test_legacy_read_failure_does_not_initialize_empty(self):
|
|
host = StorageHost()
|
|
seed_legacy(host)
|
|
host.failure = lambda action, key: action == Action.GET_PLUGIN_STORAGE
|
|
with self.assertRaises(OSError):
|
|
await open_plugin(host)
|
|
self.assertFalse(host.writes)
|
|
|
|
async def test_corrupt_legacy_data_is_not_overwritten(self):
|
|
host = StorageHost()
|
|
seed_legacy(host)
|
|
host.data["ht_messages"] = b"invalid json"
|
|
with self.assertRaises(ValueError):
|
|
await open_plugin(host)
|
|
self.assertFalse(host.writes)
|
|
|
|
async def test_interrupted_migration_can_retry_without_losing_history(self):
|
|
host = StorageHost()
|
|
expected = seed_legacy(host)
|
|
original = dict(host.data)
|
|
host.failure = lambda action, key: (
|
|
action == Action.SET_PLUGIN_STORAGE and len(host.writes) == 1
|
|
)
|
|
with self.assertRaises(OSError):
|
|
await open_plugin(host)
|
|
for key, value in original.items():
|
|
self.assertEqual(host.data[key], value)
|
|
host.failure = None
|
|
plugin = await open_plugin(host)
|
|
self.assertEqual((plugin.sessions, plugin.messages), expected)
|
|
restarted = await open_plugin(host)
|
|
self.assertEqual((restarted.sessions, restarted.messages), expected)
|
|
|
|
async def test_large_individual_session_rejected_without_truncating_history(self):
|
|
await record(self.plugin, content="previous history")
|
|
before = copy.deepcopy((self.plugin.sessions, self.plugin.messages))
|
|
stored = dict(self.host.data)
|
|
writes = len(self.host.writes)
|
|
with self.assertRaisesRegex(ValueError, "session.*limit"):
|
|
await record(self.plugin, content="界" * (4 * 1024 * 1024))
|
|
self.assertEqual((self.plugin.sessions, self.plugin.messages), before)
|
|
self.assertEqual(self.host.data, stored)
|
|
self.assertEqual(len(self.host.writes), writes)
|
|
|
|
async def test_oversized_legacy_session_fails_without_discarding_backup(self):
|
|
host = StorageHost()
|
|
seed_legacy(host, count=1, content="x" * (9 * 1024 * 1024))
|
|
original = dict(host.data)
|
|
with self.assertRaisesRegex(ValueError, "session.*limit"):
|
|
await open_plugin(host)
|
|
for key, value in original.items():
|
|
self.assertEqual(host.data[key], value)
|
|
|
|
async def test_clear_removes_legacy_and_new_data_without_resurrection(self):
|
|
host = StorageHost()
|
|
seed_legacy(host)
|
|
plugin = await open_plugin(host)
|
|
await record(plugin, "group_new")
|
|
host.data["unrelated"] = b"keep"
|
|
await plugin.clear_all()
|
|
restarted = await open_plugin(host)
|
|
self.assertEqual((restarted.sessions, restarted.messages), ({}, {}))
|
|
self.assertEqual(host.data["unrelated"], b"keep")
|
|
self.assertNotIn("ht_messages", host.data)
|
|
await record(restarted, content="after clear")
|
|
self.assertEqual((await open_plugin(host)).messages, restarted.messages)
|
|
|
|
async def test_clear_failure_is_surfaced_and_retryable(self):
|
|
await record(self.plugin)
|
|
self.host.failure = lambda action, key: (
|
|
action in {Action.DELETE_PLUGIN_STORAGE, Action.SET_PLUGIN_STORAGE}
|
|
)
|
|
with self.assertRaises(OSError):
|
|
await self.plugin.clear_all()
|
|
self.host.failure = None
|
|
await self.plugin.reconcile()
|
|
await self.plugin.clear_all()
|
|
self.assertEqual((await open_plugin(self.host)).messages, {})
|
|
|
|
async def test_null_legacy_snapshots_remain_compatible(self):
|
|
host = StorageHost()
|
|
host.data.update(ht_sessions=b"null", ht_messages=b"null")
|
|
plugin = await open_plugin(host)
|
|
self.assertEqual((plugin.sessions, plugin.messages), ({}, {}))
|
|
await record(plugin)
|
|
self.assertEqual((await open_plugin(host)).messages, plugin.messages)
|
|
|
|
async def test_invalid_legacy_bucket_prevents_migration(self):
|
|
for value in (None, "not a list", ["not a message"]):
|
|
with self.subTest(value=value):
|
|
host = StorageHost()
|
|
seed_legacy(host)
|
|
host.data["ht_messages"] = json.dumps({"group_0": value}).encode()
|
|
with self.assertRaisesRegex(ValueError, "Invalid.*storage"):
|
|
await open_plugin(host)
|
|
self.assertFalse(host.writes)
|
|
|
|
async def test_invalid_session_record_is_not_loaded(self):
|
|
await record(self.plugin)
|
|
key, raw = self.host.writes[-1]
|
|
data = json.loads(raw)
|
|
data["messages"] = "not a list"
|
|
self.host.data[key] = json.dumps(data).encode()
|
|
with self.assertRaisesRegex(ValueError, "Invalid.*storage"):
|
|
await open_plugin(self.host)
|
|
|
|
async def test_sent_message_storage_failure_warns_against_resending(self):
|
|
await record(self.plugin)
|
|
sent = []
|
|
|
|
async def send_message(**kwargs):
|
|
sent.append(kwargs)
|
|
|
|
self.plugin.send_message = send_message
|
|
self.host.failure = lambda action, key: action == Action.SET_PLUGIN_STORAGE
|
|
try:
|
|
ok, error = await self.plugin.human_send("group_1", text="manual response")
|
|
except OSError:
|
|
self.fail(
|
|
"Post-delivery storage failure must say the message was already sent"
|
|
)
|
|
self.assertFalse(ok)
|
|
self.assertIn("sent", error)
|
|
self.assertIn("storage", error)
|
|
self.assertIn("not resend", error)
|
|
self.assertEqual(len(sent), 1)
|
|
|
|
async def test_existing_300_message_retention_is_unchanged(self):
|
|
for i in range(301):
|
|
await record(self.plugin, content=str(i))
|
|
self.assertEqual(len(self.plugin.messages["group_1"]), 300)
|
|
self.assertEqual(self.plugin.messages["group_1"][0]["content"], "1")
|
|
self.assertEqual((await open_plugin(self.host)).messages, self.plugin.messages)
|
|
|
|
async def test_migration_does_not_apply_retention_to_existing_history(self):
|
|
host = StorageHost()
|
|
sessions, messages = seed_legacy(host)
|
|
messages["group_0"] = [{"content": str(i)} for i in range(400)]
|
|
host.data["ht_messages"] = json.dumps(messages).encode()
|
|
plugin = await open_plugin(host)
|
|
self.assertEqual((plugin.sessions, plugin.messages), (sessions, messages))
|
|
self.assertEqual((await open_plugin(host)).messages, messages)
|
|
|
|
async def test_clear_waits_for_inflight_record(self):
|
|
await record(self.plugin)
|
|
self.host.write_started = started = asyncio.Event()
|
|
self.host.release_write = release = asyncio.Event()
|
|
writing = asyncio.create_task(record(self.plugin, content="in flight"))
|
|
await started.wait()
|
|
clearing = asyncio.create_task(self.plugin.clear_all())
|
|
await asyncio.sleep(0.01)
|
|
self.assertFalse(clearing.done())
|
|
release.set()
|
|
await asyncio.gather(writing, clearing)
|
|
self.assertEqual((await open_plugin(self.host)).messages, {})
|
|
|
|
async def test_cancelled_write_waits_for_remote_commit(self):
|
|
for clear in (False, True):
|
|
with self.subTest(clear=clear):
|
|
host = DetachedHost()
|
|
plugin = await open_plugin(host)
|
|
await record(plugin, content="baseline")
|
|
host.pause_action = Action.SET_PLUGIN_STORAGE
|
|
writing = asyncio.create_task(record(plugin, content="cancelled"))
|
|
await host.started.wait()
|
|
writing.cancel()
|
|
await asyncio.sleep(0.01)
|
|
writing.cancel() # Repeated cancellation must not release the lock.
|
|
following = asyncio.create_task(
|
|
plugin.clear_all() if clear else record(plugin, content="next")
|
|
)
|
|
try:
|
|
await asyncio.sleep(0.01)
|
|
self.assertFalse(
|
|
following.done(), "late commit must stay serialized"
|
|
)
|
|
self.assertFalse(writing.done())
|
|
finally:
|
|
host.release.set()
|
|
await asyncio.gather(writing, following, return_exceptions=True)
|
|
await host.remote
|
|
self.assertTrue(writing.cancelled())
|
|
following.result()
|
|
restarted = await open_plugin(host)
|
|
self.assertEqual(restarted.messages, plugin.messages)
|
|
if clear:
|
|
self.assertEqual(restarted.messages, {})
|
|
else:
|
|
self.assertEqual(
|
|
[m["content"] for m in restarted.messages["group_1"]],
|
|
["baseline", "cancelled", "next"],
|
|
)
|
|
|
|
async def test_cancelled_clear_waits_for_remote_delete(self):
|
|
host = DetachedHost()
|
|
plugin = await open_plugin(host)
|
|
await record(plugin)
|
|
host.pause_action = Action.DELETE_PLUGIN_STORAGE
|
|
clearing = asyncio.create_task(plugin.clear_all())
|
|
await host.started.wait()
|
|
clearing.cancel()
|
|
await asyncio.sleep(0.01)
|
|
clearing.cancel()
|
|
following = asyncio.create_task(record(plugin, content="next"))
|
|
try:
|
|
await asyncio.sleep(0.01)
|
|
self.assertFalse(following.done())
|
|
finally:
|
|
host.release.set()
|
|
await asyncio.gather(clearing, following, return_exceptions=True)
|
|
await host.remote
|
|
self.assertTrue(clearing.cancelled())
|
|
following.result()
|
|
self.assertEqual((await open_plugin(host)).messages, plugin.messages)
|
|
self.assertEqual(len(plugin.messages["group_1"]), 1)
|
|
|
|
async def test_cancelled_migration_finishes_before_reinitialize(self):
|
|
host = DetachedHost()
|
|
expected = seed_legacy(host)
|
|
host.pause_action = Action.SET_PLUGIN_STORAGE
|
|
plugin = HumanTakeover()
|
|
plugin.config = {}
|
|
plugin.plugin_runtime_handler = host
|
|
loading = asyncio.create_task(plugin.load_state())
|
|
await host.started.wait()
|
|
loading.cancel()
|
|
await asyncio.sleep(0.01)
|
|
loading.cancel()
|
|
following = asyncio.create_task(plugin.load_state())
|
|
try:
|
|
await asyncio.sleep(0.01)
|
|
self.assertFalse(following.done())
|
|
finally:
|
|
host.release.set()
|
|
await asyncio.gather(loading, following, return_exceptions=True)
|
|
await host.remote
|
|
self.assertTrue(loading.cancelled())
|
|
following.result()
|
|
self.assertEqual((plugin.sessions, plugin.messages), expected)
|
|
|
|
async def test_ambiguous_mutations_fence_until_explicit_reconciliation(self):
|
|
for error in (TimeoutError, ConnectionError):
|
|
for operation in ("record", "clear", "migration"):
|
|
with self.subTest(error=error, operation=operation):
|
|
host = DetachedHost()
|
|
plugin = await open_plugin(host)
|
|
await record(plugin)
|
|
if operation == "migration":
|
|
host.data.pop("ht_storage_schema")
|
|
seed_legacy(host)
|
|
host.pause_action = (
|
|
Action.DELETE_PLUGIN_STORAGE
|
|
if operation == "clear"
|
|
else Action.SET_PLUGIN_STORAGE
|
|
)
|
|
host.error = error
|
|
try:
|
|
with self.assertRaises(error):
|
|
if operation == "clear":
|
|
await plugin.clear_all()
|
|
elif operation == "migration":
|
|
# Force a reload so the migration write is attempted.
|
|
await plugin.reconcile()
|
|
else:
|
|
await record(plugin, content="uncertain")
|
|
self.assertFalse(plugin._loaded)
|
|
with self.assertRaises(RuntimeError):
|
|
await record(plugin, content="must not write")
|
|
with self.assertRaises(RuntimeError):
|
|
await plugin.clear_all()
|
|
with self.assertRaisesRegex(RuntimeError, "reconcil"):
|
|
await plugin.load_state()
|
|
finally:
|
|
host.release.set()
|
|
await host.remote
|
|
# The host is now known quiescent; an ordinary read was not proof.
|
|
await plugin.reconcile()
|
|
await record(plugin, content="after reconciliation")
|
|
self.assertEqual(
|
|
(await open_plugin(host)).messages, plugin.messages
|
|
)
|
|
|
|
async def _queued_operation_cannot_cross_barrier(self, first, second, pause_action):
|
|
"""Run ``first`` so it fails only after ``second`` passed its fence check."""
|
|
host = PausableHost()
|
|
plugin = await open_plugin(host)
|
|
await record(plugin, "group_1", "baseline")
|
|
before = copy.deepcopy((plugin.sessions, plugin.messages))
|
|
before_data = dict(host.data)
|
|
host.pause_action = pause_action
|
|
host.pause_error = TimeoutError
|
|
failing = asyncio.create_task(first(plugin))
|
|
await host.paused.wait()
|
|
queued = asyncio.create_task(second(plugin))
|
|
await asyncio.sleep(0.01)
|
|
self.assertTrue(
|
|
plugin._loaded, "the barrier only exists once the writer has failed"
|
|
)
|
|
self.assertFalse(queued.done(), "the second operation must wait for the lock")
|
|
host.release.set()
|
|
with self.assertRaises(TimeoutError):
|
|
await failing
|
|
# The queued caller passed the fence check before the failure; it must
|
|
# re-check it under the lock instead of writing over an unknown outcome.
|
|
with self.assertRaisesRegex(RuntimeError, "reconcil"):
|
|
await queued
|
|
self.assertFalse(plugin._loaded)
|
|
self.assertEqual((plugin.sessions, plugin.messages), before)
|
|
self.assertEqual(host.data, before_data)
|
|
host.pause_error = None
|
|
await plugin.reconcile()
|
|
self.assertEqual((plugin.sessions, plugin.messages), before)
|
|
|
|
async def test_queued_update_cannot_cross_storage_failure_barrier(self):
|
|
await self._queued_operation_cannot_cross_barrier(
|
|
lambda plugin: record(plugin, "group_1", "not saved"),
|
|
lambda plugin: record(plugin, "group_1", "queued"),
|
|
Action.SET_PLUGIN_STORAGE,
|
|
)
|
|
|
|
async def test_queued_clear_cannot_cross_storage_failure_barrier(self):
|
|
await self._queued_operation_cannot_cross_barrier(
|
|
lambda plugin: plugin.clear_all(),
|
|
lambda plugin: record(plugin, "group_1", "queued"),
|
|
Action.DELETE_PLUGIN_STORAGE,
|
|
)
|
|
|
|
async def test_queued_update_cannot_clear_after_failed_clear(self):
|
|
await self._queued_operation_cannot_cross_barrier(
|
|
lambda plugin: plugin.clear_all(),
|
|
lambda plugin: plugin.clear_all(),
|
|
Action.DELETE_PLUGIN_STORAGE,
|
|
)
|
|
|
|
async def test_queued_loader_cannot_clear_failed_load_barrier(self):
|
|
host = PausableHost()
|
|
plugin = HumanTakeover()
|
|
plugin.config = {}
|
|
plugin.plugin_runtime_handler = host
|
|
host.pause_action = Action.GET_PLUGIN_STORAGE_KEYS
|
|
host.pause_error = ConnectionError
|
|
failing = asyncio.create_task(plugin.load_state())
|
|
await host.paused.wait()
|
|
queued = asyncio.create_task(plugin.load_state())
|
|
await asyncio.sleep(0.01)
|
|
self.assertFalse(queued.done())
|
|
host.release.set()
|
|
with self.assertRaises(ConnectionError):
|
|
await failing
|
|
# A plain load must never lift the barrier an earlier load raised, even
|
|
# though it passed the fence check before that failure.
|
|
with self.assertRaisesRegex(RuntimeError, "not initialized"):
|
|
await queued
|
|
self.assertFalse(plugin._loaded)
|
|
self.assertEqual(host.data, {})
|
|
host.pause_error = None
|
|
await plugin.reconcile()
|
|
self.assertTrue(plugin._loaded)
|
|
|
|
async def test_v2_read_failure_blocks_writes_after_failed_initialize(self):
|
|
await record(self.plugin)
|
|
self.host.failure = lambda action, key: action == Action.GET_PLUGIN_STORAGE
|
|
plugin = HumanTakeover()
|
|
plugin.config = {}
|
|
plugin.plugin_runtime_handler = self.host
|
|
with self.assertRaises(OSError):
|
|
await plugin.load_state()
|
|
self.host.failure = None
|
|
with self.assertRaisesRegex(RuntimeError, "not initialized"):
|
|
await record(plugin)
|
|
|
|
async def test_page_returns_storage_failure_not_success(self):
|
|
from components.pages.console.console import ConsolePage
|
|
from langbot_plugin.api.definition.components.page import PageRequest
|
|
|
|
await record(self.plugin)
|
|
page = ConsolePage()
|
|
page.plugin = self.plugin
|
|
self.host.failure = lambda action, key: action == Action.SET_PLUGIN_STORAGE
|
|
with patch("components.pages.console.console.logger"):
|
|
response = await page.handle_api(
|
|
PageRequest(
|
|
endpoint="/takeover", method="POST", body={"session_key": "group_1"}
|
|
)
|
|
)
|
|
self.assertIn("storage unavailable", response.error)
|
|
|
|
async def test_timeout_sweep_persists_only_changed_session(self):
|
|
await record(self.plugin, "group_1")
|
|
await record(self.plugin, "group_2")
|
|
with patch("main.time.time", return_value=1000):
|
|
await self.plugin.set_takeover("group_1", True)
|
|
self.host.writes.clear()
|
|
with patch("main.time.time", return_value=2000):
|
|
await self.plugin.expire_timeouts()
|
|
self.assertEqual(len(self.host.writes), 1)
|
|
self.assertNotIn(b"group_2", self.host.writes[0][1])
|
|
self.assertFalse(
|
|
(await open_plugin(self.host)).sessions["group_1"]["takeover"]["active"]
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|