safety layer: hard daily budget, user quotas, spend report, forgetme + privacy
This commit is contained in:
Oleksandr Kozachuk
@@ -0,0 +1,201 @@
|
||||
"""Unit coverage for FDB-014: SAF-04..09, OPS-10, PER-06."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, Mock, patch
|
||||
|
||||
from discord import TextChannel
|
||||
|
||||
from fjerkroa_bot.ai_responder import AIMessage, AIResponse
|
||||
from fjerkroa_bot.openai_responder import OpenAIResponder
|
||||
from fjerkroa_bot.persistence import SCHEMA_VERSION, PersistentStore
|
||||
from fjerkroa_bot.quota import QuotaLedger
|
||||
|
||||
from .test_bdd_envelope import envelope
|
||||
from .test_spec_ops import OpsBase
|
||||
|
||||
|
||||
def ledger_with_store(tmp, config):
|
||||
store = PersistentStore(Path(tmp) / "bot.db")
|
||||
return QuotaLedger(store, lambda: config)
|
||||
|
||||
|
||||
class TestBudgetFailClosed(unittest.IsolatedAsyncioTestCase):
|
||||
def test_spend_math_and_cutoff(self):
|
||||
"""SAF-04: spend estimate from configured prices; budget reached -> not ok."""
|
||||
config = {"daily-budget-usd": 0.01}
|
||||
ledger = QuotaLedger(None, lambda: config)
|
||||
self.assertTrue(ledger.budget_ok())
|
||||
ledger.add_tokens(20000, 0) # 20k in * $1/M = $0.02 >= $0.01
|
||||
self.assertFalse(ledger.budget_ok())
|
||||
|
||||
def test_zero_budget_is_silent(self):
|
||||
"""SAF-04: budget 0 -> fail-closed immediately."""
|
||||
ledger = QuotaLedger(None, lambda: {"daily-budget-usd": 0})
|
||||
self.assertFalse(ledger.budget_ok())
|
||||
|
||||
def test_no_budget_key_means_unlimited(self):
|
||||
"""SAF-04: without daily-budget-usd the gate stays open."""
|
||||
ledger = QuotaLedger(None, lambda: {})
|
||||
ledger.add_tokens(10_000_000, 10_000_000)
|
||||
self.assertTrue(ledger.budget_ok())
|
||||
|
||||
async def test_chat_refuses_over_budget(self):
|
||||
"""SAF-04: exhausted budget -> chat() returns None without an API call."""
|
||||
config = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5, "daily-budget-usd": 0}
|
||||
responder = OpenAIResponder(config, "chat")
|
||||
with patch("fjerkroa_bot.openai_responder.openai_chat", new_callable=AsyncMock) as chat_mock:
|
||||
answer, _ = await responder.chat([{"role": "user", "content": "hi"}], 10)
|
||||
self.assertIsNone(answer)
|
||||
chat_mock.assert_not_awaited()
|
||||
|
||||
|
||||
class TestBudgetStaffAlert(OpsBase):
|
||||
async def test_alert_once_and_silence(self):
|
||||
"""SAF-04: one staff alert per day, responder never called while exhausted."""
|
||||
self.bot.config["daily-budget-usd"] = 0
|
||||
self.bot.send_message_with_typing = AsyncMock()
|
||||
origin = MagicMock(spec=TextChannel)
|
||||
await self.bot.respond(AIMessage("alice", "hei", "chat"), origin)
|
||||
await self.bot.respond(AIMessage("alice", "hei again", "chat"), origin)
|
||||
self.bot.send_message_with_typing.assert_not_awaited()
|
||||
self.assertEqual(self.bot.staff_channel.send.await_count, 1)
|
||||
|
||||
|
||||
class TestUsageMetering(unittest.TestCase):
|
||||
def test_usage_persists_across_instances(self):
|
||||
"""SAF-05: token/image counters survive a restart via the usage table."""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
config = {"daily-budget-usd": 100}
|
||||
ledger = ledger_with_store(tmp, config)
|
||||
ledger.add_tokens(1000, 500)
|
||||
ledger.add_images(2)
|
||||
reborn = ledger_with_store(tmp, config)
|
||||
self.assertGreater(reborn.spent_usd(), 0)
|
||||
self.assertEqual(reborn.images_today(), 2)
|
||||
|
||||
|
||||
class TestChatRecordsUsage(unittest.IsolatedAsyncioTestCase):
|
||||
async def test_tokens_recorded_from_api_usage(self):
|
||||
"""SAF-05: chat() feeds prompt/completion token counts into the ledger."""
|
||||
config = {"openai-token": "t", "model": "m", "system": "s", "history-limit": 5}
|
||||
responder = OpenAIResponder(config, "chat")
|
||||
message = Mock(content=envelope(answer="x", answer_needed=True), role="assistant", tool_calls=None, refusal=None)
|
||||
usage = Mock(prompt_tokens=1234, completion_tokens=56)
|
||||
with patch("fjerkroa_bot.openai_responder.openai_chat", new_callable=AsyncMock) as chat_mock:
|
||||
chat_mock.return_value = Mock(choices=[Mock(message=message)], usage=usage)
|
||||
await responder.chat([{"role": "user", "content": "hi"}], 10)
|
||||
self.assertEqual(responder.ledger.tokens_today(), (1234, 56))
|
||||
|
||||
|
||||
class TestMessageQuota(OpsBase):
|
||||
async def test_user_over_message_cap_is_ignored(self):
|
||||
"""SAF-06: messages over user-daily-messages are dropped before the model."""
|
||||
self.bot.config["user-daily-messages"] = 2
|
||||
self.bot.send_message_with_typing = AsyncMock(return_value=AIResponse(None, False, "chat", None, None, False, False))
|
||||
origin = MagicMock(spec=TextChannel)
|
||||
for _ in range(3):
|
||||
await self.bot.respond(AIMessage("alice", "hei", "chat"), origin)
|
||||
self.assertEqual(self.bot.send_message_with_typing.await_count, 2)
|
||||
|
||||
async def test_system_user_exempt(self):
|
||||
"""SAF-06: bot-initiated (system) messages bypass the user quota."""
|
||||
self.bot.config["user-daily-messages"] = 1
|
||||
self.bot.send_message_with_typing = AsyncMock(return_value=AIResponse(None, False, "chat", None, None, False, False))
|
||||
origin = MagicMock(spec=TextChannel)
|
||||
for _ in range(3):
|
||||
await self.bot.respond(AIMessage("system", "impulse", "chat", True, False), origin)
|
||||
self.assertEqual(self.bot.send_message_with_typing.await_count, 3)
|
||||
|
||||
|
||||
class TestImageQuota(OpsBase):
|
||||
async def test_picture_stripped_over_cap(self):
|
||||
"""SAF-07: picture requests over user-daily-images are stripped, text kept."""
|
||||
self.bot.config["user-daily-images"] = 1
|
||||
self.bot.send_answer_with_typing = AsyncMock()
|
||||
origin = MagicMock(spec=TextChannel)
|
||||
for _ in range(2):
|
||||
response = AIResponse("here", True, "chat", None, "a cat", False, False)
|
||||
self.bot.send_message_with_typing = AsyncMock(return_value=response)
|
||||
await self.bot.respond(AIMessage("alice", "draw", "chat"), origin)
|
||||
first = self.bot.send_answer_with_typing.await_args_list[0].args[0]
|
||||
second = self.bot.send_answer_with_typing.await_args_list[1].args[0]
|
||||
self.assertEqual(first.picture, "a cat")
|
||||
self.assertIsNone(second.picture)
|
||||
|
||||
|
||||
class TestForgetMe(OpsBase):
|
||||
async def test_forgetme_purges_live_history_and_confirms(self):
|
||||
"""SAF-08: !forgetme removes the user's rows from live history + confirms."""
|
||||
entry = {"role": "user", "content": json.dumps({"user": "alice", "message": "secret", "channel": "chat"})}
|
||||
other = {"role": "user", "content": json.dumps({"user": "bob", "message": "stays", "channel": "chat"})}
|
||||
self.bot.airesponder.history = [dict(entry), dict(other)]
|
||||
message = self.public_msg("!forgetme")
|
||||
message.author.name = "alice"
|
||||
await self.bot.on_message(message)
|
||||
contents = [item["content"] for item in self.bot.airesponder.history]
|
||||
self.assertFalse(any('"alice"' in content for content in contents))
|
||||
self.assertTrue(any('"bob"' in content for content in contents))
|
||||
message.channel.send.assert_awaited_once()
|
||||
|
||||
def test_store_purge_by_user(self):
|
||||
"""SAF-08: the store deletes persisted rows containing the user's messages."""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
store = PersistentStore(Path(tmp) / "bot.db")
|
||||
store.save_history(
|
||||
"chat",
|
||||
[
|
||||
{"role": "user", "content": json.dumps({"user": "alice", "message": "x"})},
|
||||
{"role": "user", "content": json.dumps({"user": "bob", "message": "y"})},
|
||||
],
|
||||
)
|
||||
store.delete_history_of_user("alice")
|
||||
remaining = store.load_history("chat")
|
||||
self.assertEqual(len(remaining), 1)
|
||||
self.assertIn('"bob"', remaining[0]["content"])
|
||||
|
||||
|
||||
class TestPrivacyNotice(OpsBase):
|
||||
async def test_privacy_answers_even_when_paused(self):
|
||||
"""SAF-09: !privacy answers with the notice, also while paused."""
|
||||
self.bot.replies_enabled = False
|
||||
self.bot.config["privacy-notice"] = "We store recent messages. Use !forgetme."
|
||||
message = self.public_msg("!privacy")
|
||||
await self.bot.on_message(message)
|
||||
message.channel.send.assert_awaited_once()
|
||||
self.assertIn("!forgetme", message.channel.send.await_args.args[0])
|
||||
|
||||
|
||||
class TestSpendCommand(OpsBase):
|
||||
async def test_spend_report(self):
|
||||
"""OPS-10: !bot spend reports estimated USD + counters + budget."""
|
||||
await self.bot.on_message(self.staff_msg("!bot spend"))
|
||||
self.bot.staff_channel.send.assert_awaited()
|
||||
text = self.bot.staff_channel.send.await_args.args[0]
|
||||
self.assertIn("$", text)
|
||||
self.assertIn("tokens", text)
|
||||
|
||||
|
||||
class TestSchemaMigration(unittest.TestCase):
|
||||
def test_v1_database_upgrades_to_current(self):
|
||||
"""PER-06: v1 db gains the usage table, keeps rows, bumps user_version."""
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
db_path = Path(tmp) / "bot.db"
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.execute("CREATE TABLE history (id INTEGER PRIMARY KEY, channel TEXT NOT NULL, role TEXT NOT NULL, content TEXT NOT NULL)")
|
||||
conn.execute("CREATE TABLE memory (channel TEXT PRIMARY KEY, content TEXT NOT NULL)")
|
||||
conn.execute("INSERT INTO history (channel, role, content) VALUES ('chat', 'user', 'kept')")
|
||||
conn.execute("PRAGMA user_version = 1")
|
||||
conn.commit()
|
||||
conn.close()
|
||||
store = PersistentStore(db_path)
|
||||
self.assertEqual(store.load_history("chat"), [{"role": "user", "content": "kept"}])
|
||||
store.usage_add("2026-07-13", "tokens-in", 5)
|
||||
check = sqlite3.connect(db_path)
|
||||
try:
|
||||
self.assertEqual(check.execute("PRAGMA user_version").fetchone()[0], SCHEMA_VERSION)
|
||||
finally:
|
||||
check.close()
|
||||
Reference in New Issue
Block a user