107 lines
4.3 KiB
Python
107 lines
4.3 KiB
Python
import sqlite3
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from aiogram import Dispatcher
|
|
from aiogram.types import Update
|
|
|
|
import bot
|
|
import raw_updates
|
|
|
|
|
|
class RawUpdateStorageTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.db_path = str(Path(self.temp_dir.name) / "updates.sqlite3")
|
|
raw_updates.init_raw_updates_db(self.db_path)
|
|
|
|
def tearDown(self):
|
|
self.temp_dir.cleanup()
|
|
|
|
def test_save_and_read_update_by_id(self):
|
|
payload = {"update_id": 456, "message": {"message_id": 8, "text": "Привет"}}
|
|
raw_updates.save_raw_update(456, payload, self.db_path)
|
|
record = raw_updates.get_raw_update(456, self.db_path)
|
|
self.assertIsNotNone(record)
|
|
self.assertEqual(record["event_type"], "message")
|
|
self.assertEqual(record["payload"], payload)
|
|
self.assertIsNone(raw_updates.get_raw_update(999, self.db_path))
|
|
|
|
def test_sensitive_fields_are_redacted(self):
|
|
raw_updates.save_raw_update(
|
|
457,
|
|
{"update_id": 457, "message": {"text": "Bearer abc123", "token": "private"}},
|
|
self.db_path,
|
|
)
|
|
record = raw_updates.get_raw_update(457, self.db_path)
|
|
assert record is not None
|
|
self.assertEqual(record["payload"]["message"]["text"], "Bearer [REDACTED]")
|
|
self.assertEqual(record["payload"]["message"]["token"], "[REDACTED]")
|
|
|
|
def test_records_older_than_retention_are_removed(self):
|
|
raw_updates.save_raw_update(458, {"update_id": 458, "message": {}}, self.db_path)
|
|
with sqlite3.connect(self.db_path) as connection:
|
|
connection.execute(
|
|
"UPDATE raw_updates SET received_at = datetime('now', '-31 days') WHERE update_id = 458"
|
|
)
|
|
raw_updates.init_raw_updates_db(self.db_path)
|
|
self.assertIsNone(raw_updates.get_raw_update(458, self.db_path))
|
|
|
|
|
|
class RawUpdateMiddlewareTests(unittest.IsolatedAsyncioTestCase):
|
|
async def test_update_is_persisted_before_handler(self):
|
|
update = Update.model_validate({
|
|
"update_id": 456,
|
|
"message": {
|
|
"message_id": 8,
|
|
"date": 1,
|
|
"chat": {"id": 42, "type": "private"},
|
|
"from": {"id": 42, "is_bot": False, "first_name": "Test"},
|
|
"text": "Привет",
|
|
},
|
|
})
|
|
temp_dir = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temp_dir.cleanup)
|
|
db_path = str(Path(temp_dir.name) / "updates.sqlite3")
|
|
raw_updates.init_raw_updates_db(db_path)
|
|
middleware = raw_updates.RawUpdateMiddleware()
|
|
|
|
async def handler(event, data):
|
|
record = raw_updates.get_raw_update(456, db_path)
|
|
assert record is not None
|
|
self.assertEqual(record["payload"]["message"]["text"], "Привет")
|
|
return "handled"
|
|
|
|
with (
|
|
patch.object(raw_updates, "RAW_UPDATES_DB_PATH", db_path),
|
|
patch.object(raw_updates, "save_raw_update", wraps=raw_updates.save_raw_update) as save,
|
|
):
|
|
result = await middleware(handler, update, {})
|
|
self.assertEqual(result, "handled")
|
|
save.assert_called_once()
|
|
|
|
async def test_persistence_failure_does_not_block_update(self):
|
|
middleware = raw_updates.RawUpdateMiddleware()
|
|
update = Update(update_id=459)
|
|
handler = AsyncMock(return_value="handled")
|
|
with patch.object(raw_updates, "save_raw_update", side_effect=OSError("disk full")):
|
|
self.assertEqual(await middleware(handler, update, {}), "handled")
|
|
handler.assert_awaited_once()
|
|
|
|
|
|
class RawUpdateModeTests(unittest.TestCase):
|
|
def test_middleware_is_registered_only_for_development(self):
|
|
dispatcher = Dispatcher()
|
|
with patch.object(bot, "DEPLOY_ENV", "production"), patch.object(
|
|
dispatcher.update, "outer_middleware"
|
|
) as register:
|
|
self.assertFalse(bot.configure_raw_update_logging(dispatcher))
|
|
register.assert_not_called()
|
|
with patch.object(bot, "DEPLOY_ENV", "development"), patch.object(
|
|
dispatcher.update, "outer_middleware"
|
|
) as register:
|
|
self.assertTrue(bot.configure_raw_update_logging(dispatcher))
|
|
register.assert_called_once()
|