import json import tempfile import unittest from datetime import datetime, timedelta, timezone from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, patch from bot.services.panel_api_service import PanelApiService from bot.services.subscription_service import SubscriptionService from bot.services.tariff_worker import TariffTrafficWorker from config.settings import Settings def _tariffs_config_payload() -> dict: return { "default_tariff": "standard", "tariffs": [ { "key": "standard", "names": {"ru": "Стандарт"}, "descriptions": {"ru": "Base"}, "squad_uuids": ["squad-1"], "billing_model": "period", "monthly_gb": 500, "prices_rub": {"1": 150}, "prices_stars": {"1": 0}, "enabled_periods": [1], "enabled": True, } ], } class TariffWorkerTests(unittest.IsolatedAsyncioTestCase): async def test_period_tariff_uses_panel_month_strategy_without_resetting(self): with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(_tariffs_config_payload()), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), ) panel_service = AsyncMock(spec=PanelApiService) panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) panel_service.reset_user_traffic = AsyncMock(return_value=True) panel_service.add_users_to_internal_squad = AsyncMock(return_value=True) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) sub = SimpleNamespace( subscription_id=1, user_id=123, panel_user_uuid="panel-uuid", end_date=datetime.now(timezone.utc) + timedelta(days=10), traffic_limit_bytes=500 * (1024**3), topup_balance_bytes=0, is_throttled=False, status_from_panel="ACTIVE", ) tariff = settings.tariffs_config.require("standard") await worker._ensure_period_reset_strategy(sub, tariff, sub.traffic_limit_bytes, "NO_RESET") panel_service.update_user_details_on_panel.assert_awaited_once() panel_service.reset_user_traffic.assert_not_awaited() update_payload = panel_service.update_user_details_on_panel.await_args.args[1] self.assertEqual(update_payload["trafficLimitStrategy"], "MONTH") self.assertEqual(update_payload["trafficLimitBytes"], sub.traffic_limit_bytes) self.assertNotIn("status", update_payload) async def test_limit_reached_does_not_remove_user_from_squad(self): with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(_tariffs_config_payload()), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), TARIFF_TRAFFIC_WARNING_LEVELS="101", ) panel_service = AsyncMock(spec=PanelApiService) panel_service.remove_users_from_internal_squad = AsyncMock(return_value=True) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) sub = SimpleNamespace( subscription_id=1, user_id=123, panel_user_uuid="panel-uuid", traffic_limit_bytes=100, traffic_used_bytes=100, is_throttled=False, status_from_panel="ACTIVE", ) tariff = settings.tariffs_config.require("standard") with patch("bot.services.tariff_worker.tariff_dal.get_warning", new=AsyncMock(return_value=True)): await worker._maybe_warn_or_throttle( AsyncMock(), sub, tariff, used=100, limit=100, warning_period_start=datetime.now(timezone.utc), ) panel_service.remove_users_from_internal_squad.assert_not_awaited() self.assertFalse(sub.is_throttled) async def test_premium_limit_removes_only_premium_squad(self): payload = _tariffs_config_payload() payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"] payload["tariffs"][0]["premium_monthly_gb"] = 1 with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(payload), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), ) panel_service = AsyncMock(spec=PanelApiService) panel_service.get_internal_squad_accessible_nodes = AsyncMock( return_value=[{"uuid": "node-1", "name": "Premium"}] ) panel_service.get_node_users_bandwidth_stats = AsyncMock( return_value={ "topUsers": [ { "username": "tg_123", "total": 2 * (1024**3), } ] } ) panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) sub = SimpleNamespace( subscription_id=1, user_id=123, panel_user_uuid="panel-uuid", premium_baseline_bytes=1 * (1024**3), premium_topup_balance_bytes=0, premium_topup_used_bytes=0, premium_used_bytes=0, premium_is_limited=False, premium_period_start_at=None, ) tariff = settings.tariffs_config.require("standard") with patch("bot.services.tariff_worker.tariff_dal.get_warning", new=AsyncMock(return_value=True)): await worker._sync_premium_squad_limit( AsyncMock(), sub, tariff, datetime.now(timezone.utc), panel_username="tg_123", ) self.assertTrue(sub.premium_is_limited) panel_service.update_user_details_on_panel.assert_awaited_once() payload = panel_service.update_user_details_on_panel.await_args.args[1] self.assertEqual(payload["activeInternalSquads"], ["squad-1"]) async def test_premium_topup_balance_carries_over_and_is_spent_only_above_monthly_limit(self): payload = _tariffs_config_payload() payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"] payload["tariffs"][0]["premium_monthly_gb"] = 1 with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(payload), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), TARIFF_TRAFFIC_WARNING_LEVELS="101", ) panel_service = AsyncMock(spec=PanelApiService) panel_service.get_internal_squad_accessible_nodes = AsyncMock(return_value=[{"uuid": "node-1"}]) panel_service.get_node_users_bandwidth_stats = AsyncMock( return_value={ "topUsers": [ { "username": "tg_123", "total": int(1.5 * (1024**3)), } ] } ) panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) now = datetime(2026, 5, 9, tzinfo=timezone.utc) sub = SimpleNamespace( subscription_id=1, user_id=123, panel_user_uuid="panel-uuid", premium_baseline_bytes=1 * (1024**3), premium_topup_balance_bytes=2 * (1024**3), premium_topup_used_bytes=0, premium_used_bytes=0, premium_is_limited=False, premium_period_start_at=datetime(2026, 5, 1, tzinfo=timezone.utc), ) tariff = settings.tariffs_config.require("standard") await worker._sync_premium_squad_limit( AsyncMock(), sub, tariff, now, panel_username="tg_123" ) self.assertEqual(sub.premium_topup_balance_bytes, int(1.5 * (1024**3))) self.assertEqual(sub.premium_topup_used_bytes, int(0.5 * (1024**3))) self.assertFalse(sub.premium_is_limited) panel_service.get_node_users_bandwidth_stats = AsyncMock( return_value={ "topUsers": [ { "username": "tg_123", "total": int(0.1 * (1024**3)), } ] } ) next_month = datetime(2026, 6, 2, tzinfo=timezone.utc) await worker._sync_premium_squad_limit( AsyncMock(), sub, tariff, next_month, panel_username="tg_123" ) self.assertEqual(sub.premium_topup_balance_bytes, int(1.5 * (1024**3))) self.assertEqual(sub.premium_topup_used_bytes, 0) self.assertEqual(sub.premium_period_start_at, datetime(2026, 6, 1, tzinfo=timezone.utc)) async def test_premium_unlimited_override_never_throttles(self): payload = _tariffs_config_payload() payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"] payload["tariffs"][0]["premium_monthly_gb"] = 1 with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(payload), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), TARIFF_TRAFFIC_WARNING_LEVELS="101", ) panel_service = AsyncMock(spec=PanelApiService) panel_service.get_internal_squad_accessible_nodes = AsyncMock( return_value=[{"uuid": "node-1"}] ) panel_service.get_node_users_bandwidth_stats = AsyncMock( return_value={ "topUsers": [ {"username": "tg_42", "total": 50 * (1024**3)}, ] } ) panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) sub = SimpleNamespace( subscription_id=1, user_id=42, panel_user_uuid="panel-uuid", premium_baseline_bytes=1 * (1024**3), premium_topup_balance_bytes=0, premium_topup_used_bytes=0, premium_used_bytes=0, premium_is_limited=False, premium_period_start_at=None, premium_unlimited_override=True, premium_bonus_bytes=0, ) tariff = settings.tariffs_config.require("standard") await worker._sync_premium_squad_limit( AsyncMock(), sub, tariff, datetime.now(timezone.utc), panel_username="tg_42" ) self.assertFalse(sub.premium_is_limited) self.assertEqual(int(sub.premium_used_bytes), 50 * (1024**3)) payload_sent = panel_service.update_user_details_on_panel.await_args.args[1] self.assertIn("premium-squad", payload_sent["activeInternalSquads"]) async def test_premium_bonus_extends_limit(self): payload = _tariffs_config_payload() payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"] payload["tariffs"][0]["premium_monthly_gb"] = 1 with tempfile.TemporaryDirectory() as tmpdir: config_path = Path(tmpdir) / "tariffs.json" config_path.write_text(json.dumps(payload), encoding="utf-8") settings = Settings( _env_file=None, BOT_TOKEN="token", POSTGRES_USER="app_user", POSTGRES_PASSWORD="app_password", TARIFFS_CONFIG_PATH=str(config_path), TARIFF_TRAFFIC_WARNING_LEVELS="101", ) panel_service = AsyncMock(spec=PanelApiService) panel_service.get_internal_squad_accessible_nodes = AsyncMock( return_value=[{"uuid": "node-1"}] ) # Used 4 GB > tariff baseline 1 GB, but admin granted +10 GB bonus. panel_service.get_node_users_bandwidth_stats = AsyncMock( return_value={ "topUsers": [ {"username": "tg_77", "total": 4 * (1024**3)}, ] } ) panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}}) subscription_service = SubscriptionService(settings, panel_service) worker = TariffTrafficWorker( settings=settings, session_factory=SimpleNamespace(), panel_service=panel_service, subscription_service=subscription_service, ) sub = SimpleNamespace( subscription_id=1, user_id=77, panel_user_uuid="panel-uuid", premium_baseline_bytes=1 * (1024**3), premium_topup_balance_bytes=0, premium_topup_used_bytes=0, premium_used_bytes=0, premium_is_limited=False, premium_period_start_at=None, premium_unlimited_override=False, premium_bonus_bytes=10 * (1024**3), ) tariff = settings.tariffs_config.require("standard") await worker._sync_premium_squad_limit( AsyncMock(), sub, tariff, datetime.now(timezone.utc), panel_username="tg_77" ) # 4 GB used vs 1 GB baseline + 10 GB bonus = 11 GB limit → not limited. self.assertFalse(sub.premium_is_limited) self.assertEqual(int(sub.premium_used_bytes), 4 * (1024**3))