121 lines
5.3 KiB
Python
121 lines
5.3 KiB
Python
import asyncio
|
|
from datetime import datetime, timezone
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
from aiogram import Bot, Dispatcher, Router, F
|
|
from aiogram.fsm.storage.memory import MemoryStorage, SimpleEventIsolation
|
|
from aiogram.fsm.context import FSMContext
|
|
from aiogram.types import Message, Chat, User as TelegramUser, CallbackQuery, Update
|
|
from sqlalchemy import select
|
|
|
|
from src.core.access import is_admin, is_staff
|
|
from src.core.database import async_session_maker
|
|
from src.core.models import User, Winner
|
|
from src.core.services import UserService
|
|
from src.middlewares.access import AccessMiddleware
|
|
from src.controllers.bot_controller import BotController
|
|
|
|
|
|
async def test_assigned_admin_and_cashier_permissions_update_immediately():
|
|
async with async_session_maker() as session:
|
|
admin = await UserService.get_or_create_user(session, 301)
|
|
admin.is_admin = True
|
|
cashier = await UserService.get_or_create_user(session, 302)
|
|
cashier.is_cashier = True
|
|
await session.commit()
|
|
assert await is_admin(301)
|
|
assert await is_staff(302) and not await is_admin(302)
|
|
admin.is_admin = False
|
|
await session.commit()
|
|
assert not await is_admin(301)
|
|
|
|
|
|
@pytest.mark.parametrize("role,allow_cashier,permitted", [("user", False, False), ("user", True, False),
|
|
("cashier", False, False), ("cashier", True, True), ("admin", False, True)])
|
|
async def test_every_staff_fsm_step_checks_permissions(role, allow_cashier, permitted):
|
|
async with async_session_maker() as session:
|
|
user = await UserService.get_or_create_user(session, 303)
|
|
user.is_admin = role == "admin"
|
|
user.is_cashier = role == "cashier"
|
|
await session.commit()
|
|
handler = AsyncMock()
|
|
event = SimpleNamespace(from_user=SimpleNamespace(id=303), answer=AsyncMock())
|
|
await AccessMiddleware(allow_cashier=allow_cashier)(handler, event, {})
|
|
assert handler.await_count == int(permitted)
|
|
|
|
|
|
async def test_callback_main_menu_uses_clicker_identity():
|
|
from main import back_to_main_handler
|
|
import main
|
|
from contextlib import asynccontextmanager
|
|
controller = SimpleNamespace(handle_start=AsyncMock())
|
|
@asynccontextmanager
|
|
async def get_controller():
|
|
yield controller
|
|
old = main.get_controller
|
|
main.get_controller = get_controller
|
|
try:
|
|
actor = TelegramUser(id=304, is_bot=False, first_name="Human")
|
|
event = SimpleNamespace(from_user=actor, message=object(), answer=AsyncMock())
|
|
state = SimpleNamespace(clear=AsyncMock())
|
|
await back_to_main_handler(event, state)
|
|
controller.handle_start.assert_awaited_once_with(event.message, actor=actor)
|
|
state.clear.assert_awaited_once()
|
|
finally:
|
|
main.get_controller = old
|
|
|
|
|
|
async def test_slow_user_does_not_block_other_user_and_same_user_is_serialized():
|
|
storage = MemoryStorage()
|
|
dispatcher = Dispatcher(storage=storage, events_isolation=SimpleEventIsolation())
|
|
router = Router()
|
|
bot = Bot("123456:TEST_TOKEN_FOR_ISOLATED_TESTS")
|
|
first_started, release_first, other_finished = asyncio.Event(), asyncio.Event(), asyncio.Event()
|
|
seen = []
|
|
@router.message(F.text)
|
|
async def receive(message: Message, state: FSMContext):
|
|
if message.text == "slow":
|
|
first_started.set()
|
|
await release_first.wait()
|
|
data = await state.get_data()
|
|
seen.append((message.from_user.id, data.get("count", 0)))
|
|
await state.update_data(count=data.get("count", 0) + 1)
|
|
if message.from_user.id == 402:
|
|
other_finished.set()
|
|
dispatcher.include_router(router)
|
|
def update(number, user_id, text):
|
|
return Update(update_id=number, message=Message(message_id=number, date=datetime.now(timezone.utc),
|
|
chat=Chat(id=user_id, type="private"), from_user=TelegramUser(id=user_id, is_bot=False, first_name="Actor"), text=text))
|
|
slow = asyncio.create_task(dispatcher.feed_update(bot, update(1, 401, "slow")))
|
|
await asyncio.wait_for(first_started.wait(), 1)
|
|
same = asyncio.create_task(dispatcher.feed_update(bot, update(2, 401, "next")))
|
|
other = asyncio.create_task(dispatcher.feed_update(bot, update(3, 402, "fast")))
|
|
try:
|
|
await asyncio.wait_for(other_finished.wait(), 1)
|
|
assert not same.done()
|
|
finally:
|
|
release_first.set()
|
|
await asyncio.gather(slow, same, other)
|
|
await bot.session.close()
|
|
await dispatcher.fsm.close()
|
|
assert seen == [(402, 0), (401, 0), (401, 1)]
|
|
|
|
|
|
async def test_non_owner_cannot_confirm_user_only_prize():
|
|
from src.handlers.redraw_handlers import confirm_winner_callback
|
|
from src.core.services import LotteryService
|
|
async with async_session_maker() as session:
|
|
owner = await UserService.get_or_create_user(session, 601)
|
|
draw = await LotteryService.create_lottery(session, "Draw", "", ["Prize"], owner.id)
|
|
winner = Winner(lottery_id=draw.id, user_id=owner.id, place=1)
|
|
session.add(winner)
|
|
await session.commit()
|
|
winner_id = winner.id
|
|
callback = SimpleNamespace(data=f"confirm_win_{winner_id}", from_user=SimpleNamespace(id=602), answer=AsyncMock())
|
|
await confirm_winner_callback(callback)
|
|
async with async_session_maker() as session:
|
|
assert not (await session.get(Winner, winner_id)).is_claimed
|
|
callback.answer.assert_awaited_once()
|