Oleksandr Kozachuk
202 lines
9.7 KiB
Python
202 lines
9.7 KiB
Python
"""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()
|