diff --git a/bot/services/subscription_service.py b/bot/services/subscription_service.py index 0a22226..7f4f239 100644 --- a/bot/services/subscription_service.py +++ b/bot/services/subscription_service.py @@ -8,7 +8,7 @@ from bot.middlewares.i18n import JsonI18n from db.dal import user_dal, subscription_dal, promo_code_dal, payment_dal, user_billing_dal, tariff_dal from config.tariffs_config import Tariff -from bot.utils.date_utils import add_months, month_start +from bot.utils.date_utils import add_months from bot.utils.config_link import prepare_config_links from db.models import User, Subscription @@ -825,7 +825,7 @@ class SubscriptionService: if target.billing_model == "period": update_data["tier_baseline_bytes"] = target.monthly_bytes update_data["traffic_limit_bytes"] = target.monthly_bytes + int(sub.topup_balance_bytes or 0) - update_data["period_start_at"] = month_start() + update_data["period_start_at"] = None update_data["effective_monthly_price_rub"] = target.period_price(1, "rub") or target.min_period_price_rub() if mode == "recalc_days" and options.get("recalc_days") is not None: update_data["end_date"] = now + timedelta(days=int(options["recalc_days"])) @@ -858,7 +858,7 @@ class SubscriptionService: expire_at=updated.end_date, status="ACTIVE", traffic_limit_bytes=updated.traffic_limit_bytes, - traffic_limit_strategy="NO_RESET" if target.billing_model == "traffic" else self.settings.USER_TRAFFIC_STRATEGY, + traffic_limit_strategy="NO_RESET" if target.billing_model == "traffic" else "MONTH", ) panel_payload["activeInternalSquads"] = target.squad_uuids panel_payload.update(self._panel_identity_payload_for_user(db_user)) @@ -1077,11 +1077,6 @@ class SubscriptionService: topup_balance_bytes = int(getattr(current_active_sub, "topup_balance_bytes", 0) or 0) tier_baseline_bytes = tariff.monthly_bytes if tariff else self.settings.user_traffic_limit_bytes - period_start_at = month_start() if tariff else ( - datetime.now(timezone.utc) - if starts_after_lapse or not current_active_sub or not getattr(current_active_sub, "period_start_at", None) - else current_active_sub.period_start_at - ) effective_monthly_price = float(payment_amount) / max(1, months_int) traffic_limit_bytes = self._traffic_limit_for_period_tariff(tariff, topup_balance_bytes) sub_payload = { @@ -1100,7 +1095,7 @@ class SubscriptionService: "tariff_key": tariff.key if tariff else None, "tier_baseline_bytes": tier_baseline_bytes, "topup_balance_bytes": topup_balance_bytes, - "period_start_at": period_start_at, + "period_start_at": None, "is_throttled": False, "effective_monthly_price_rub": effective_monthly_price, } @@ -1120,6 +1115,7 @@ class SubscriptionService: expire_at=final_end_date, status="ACTIVE", traffic_limit_bytes=traffic_limit_bytes, + traffic_limit_strategy="MONTH" if tariff else self.settings.USER_TRAFFIC_STRATEGY, ) if tariff: panel_update_payload["activeInternalSquads"] = tariff.squad_uuids @@ -1381,7 +1377,7 @@ class SubscriptionService: except Exception: tariff = None billing_model_display = tariff.billing_model if tariff else ("traffic" if getattr(self.settings, "traffic_sale_mode", False) else "period") - traffic_limit_strategy = "MONTH" if billing_model_display == "period" else panel_traffic_strategy + traffic_limit_strategy = panel_traffic_strategy return { "user_id": panel_user_data.get("uuid"), diff --git a/bot/services/tariff_worker.py b/bot/services/tariff_worker.py index 1896cf3..dd9219b 100644 --- a/bot/services/tariff_worker.py +++ b/bot/services/tariff_worker.py @@ -58,6 +58,7 @@ class TariffTrafficWorker: async def traffic_period_tick(self, session: AsyncSession) -> None: now = datetime.now(timezone.utc) + warning_period_start = month_start(now) result = await session.execute( select(Subscription).where( Subscription.is_active == True, @@ -71,61 +72,42 @@ class TariffTrafficWorker: except Exception: continue panel_data = await self.panel_service.get_user_by_uuid(sub.panel_user_uuid, log_response=False) or {} - used, limit, _ = self.subscription_service._extract_panel_traffic_details(panel_data) + used, limit, panel_strategy = self.subscription_service._extract_panel_traffic_details(panel_data) if used is not None and used != sub.traffic_used_bytes: sub.traffic_used_bytes = used if limit is not None and limit != sub.traffic_limit_bytes: sub.traffic_limit_bytes = limit if tariff.billing_model == "period": - reset_happened = await self._maybe_reset_period(session, sub, tariff, used) - if reset_happened: - used = 0 - sub.traffic_used_bytes = 0 - await self._maybe_warn_or_throttle(session, sub, tariff, used, limit) - - async def _maybe_reset_period(self, session: AsyncSession, sub: Subscription, tariff, used: Optional[int]) -> bool: - now = datetime.now(timezone.utc) - current_month_start = month_start(now) - if not sub.period_start_at: - await subscription_dal.update_subscription( + await self._ensure_period_reset_strategy(sub, tariff, limit, panel_strategy) + await self._maybe_warn_or_throttle( session, - sub.subscription_id, - {"period_start_at": current_month_start}, + sub, + tariff, + used, + limit, + warning_period_start=warning_period_start if tariff.billing_model == "period" else None, ) - sub.period_start_at = current_month_start - return False - stored_month_start = month_start(sub.period_start_at) - if stored_month_start == current_month_start or sub.end_date <= now: - return False - await self.panel_service.reset_user_traffic(sub.panel_user_uuid) - restore_throttled = bool(sub.is_throttled) - restore_succeeded = True - if restore_throttled: - for squad_uuid in tariff.squad_uuids: - restore_succeeded = await self.panel_service.add_users_to_internal_squad( - squad_uuid, - [sub.panel_user_uuid], - ) and restore_succeeded - should_clear_throttle = restore_throttled and restore_succeeded - await subscription_dal.update_subscription( - session, - sub.subscription_id, - { - "period_start_at": current_month_start, - "traffic_used_bytes": 0, - "is_throttled": False if should_clear_throttle else sub.is_throttled, - "status_from_panel": "ACTIVE" if should_clear_throttle else sub.status_from_panel, - }, + async def _ensure_period_reset_strategy( + self, + sub: Subscription, + tariff, + limit: Optional[int], + panel_strategy: Optional[str], + ) -> None: + if str(panel_strategy or "").upper() == "MONTH": + return + traffic_limit_bytes = int(limit or sub.traffic_limit_bytes or (tariff.monthly_bytes + int(sub.topup_balance_bytes or 0))) + payload = self.subscription_service._build_panel_update_payload( + panel_user_uuid=sub.panel_user_uuid, + expire_at=sub.end_date, + status="ACTIVE", + traffic_limit_bytes=traffic_limit_bytes, + traffic_limit_strategy="MONTH", ) - sub.period_start_at = current_month_start - sub.traffic_used_bytes = 0 - if should_clear_throttle: - sub.is_throttled = False - sub.status_from_panel = "ACTIVE" - await tariff_dal.clear_period_warnings(session, sub.subscription_id) - return True + payload["activeInternalSquads"] = tariff.squad_uuids + await self.panel_service.update_user_details_on_panel(sub.panel_user_uuid, payload, log_response=False) async def _maybe_warn_or_throttle( self, @@ -134,6 +116,8 @@ class TariffTrafficWorker: tariff, used: Optional[int], limit: Optional[int], + *, + warning_period_start: Optional[datetime] = None, ) -> None: used_val = int(used or sub.traffic_used_bytes or 0) limit_val = int(limit or sub.traffic_limit_bytes or 0) @@ -146,7 +130,7 @@ class TariffTrafficWorker: warning = await tariff_dal.get_warning( session, subscription_id=sub.subscription_id, - period_start_at=sub.period_start_at if tariff.billing_model == "period" else None, + period_start_at=warning_period_start if tariff.billing_model == "period" else None, level=level, traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None, ) @@ -155,7 +139,7 @@ class TariffTrafficWorker: await tariff_dal.create_warning( session, subscription_id=sub.subscription_id, - period_start_at=sub.period_start_at if tariff.billing_model == "period" else None, + period_start_at=warning_period_start if tariff.billing_model == "period" else None, level=level, traffic_limit_bytes=limit_val if tariff.billing_model == "traffic" else None, ) diff --git a/db/database_setup.py b/db/database_setup.py index 813a389..9a94bfc 100644 --- a/db/database_setup.py +++ b/db/database_setup.py @@ -1,13 +1,9 @@ import logging - -from sqlalchemy import select from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker from sqlalchemy.orm import sessionmaker -from bot.utils.date_utils import month_start from config.settings import Settings -from db.dal import subscription_dal -from db.models import Base, Subscription +from db.models import Base from .migrator import run_database_migrations async_engine = None @@ -105,19 +101,7 @@ async def init_db(settings: Settings, session_factory: sessionmaker): tariff_key = COALESCE(s.tariff_key, :tariff_key), tier_baseline_bytes = COALESCE(s.tier_baseline_bytes, s.traffic_limit_bytes, :baseline), topup_balance_bytes = COALESCE(s.topup_balance_bytes, 0), - period_start_at = COALESCE( - s.period_start_at, - ( - SELECT p.created_at - FROM payments p - WHERE p.user_id = s.user_id - AND p.status = 'succeeded' - ORDER BY p.created_at DESC - LIMIT 1 - ), - s.start_date, - NOW() - ), + period_start_at = NULL, effective_monthly_price_rub = COALESCE( s.effective_monthly_price_rub, ( @@ -141,49 +125,16 @@ async def init_db(settings: Settings, session_factory: sessionmaker): "default_price": default_price, }, ) - month_anchor = month_start() - active_subs = await session.execute( - select(Subscription).where(Subscription.is_active == True) + await session.execute( + text( + """ + UPDATE subscriptions + SET period_start_at = NULL + WHERE is_active = TRUE + AND tariff_key IS NOT NULL + """ + ) ) - for sub in active_subs.scalars().all(): - tariff_key = sub.tariff_key or default_tariff.key - try: - tariff = settings.tariffs_config.require(tariff_key) - except Exception: - continue - - update_data = { - "tariff_key": tariff.key, - "topup_balance_bytes": int(sub.topup_balance_bytes or 0), - } - if tariff.billing_model == "period": - baseline_source = ( - sub.tier_baseline_bytes - if sub.tier_baseline_bytes is not None - else (sub.traffic_limit_bytes if sub.traffic_limit_bytes is not None else tariff.monthly_bytes) - ) - baseline = int(baseline_source or 0) - update_data.update( - { - "tier_baseline_bytes": baseline, - "traffic_limit_bytes": baseline + int(sub.topup_balance_bytes or 0), - "period_start_at": month_anchor, - "effective_monthly_price_rub": ( - sub.effective_monthly_price_rub - if sub.effective_monthly_price_rub is not None - else default_price - ), - } - ) - else: - update_data.update( - { - "tier_baseline_bytes": 0, - "period_start_at": None, - "effective_monthly_price_rub": None, - } - ) - await subscription_dal.update_subscription(session, sub.subscription_id, update_data) await session.commit() except Exception: await session.rollback() diff --git a/docs/tariffs.md b/docs/tariffs.md index 064dcb4..7da068f 100644 --- a/docs/tariffs.md +++ b/docs/tariffs.md @@ -26,9 +26,10 @@ Period-тариф продает доступ на срок и лимит тра - `monthly_gb` превращается в `tier_baseline_bytes`. - Докупленные пакеты хранятся в `topup_balance_bytes`. - В Remnawave пушится `trafficLimitBytes = tier_baseline_bytes + topup_balance_bytes`. -- Сброс происходит в начале календарного месяца для всех одинаково. -- Если покупка или продление были в середине месяца, reset всё равно придёт на ближайший первый день следующего месяца. -- Панель по-прежнему считает `usedTrafficBytes`, а бот только синхронизирует месячный reset и лимиты. +- Для period-тарифов бот выставляет `trafficLimitStrategy = MONTH`, а дальнейший reset делает сама панель. +- Дата сброса больше не считается в боте. +- Если покупка или продление были в середине месяца, сброс всё равно произойдёт по правилам панели для `MONTH`. +- Бот только меняет лимиты в GB и следит за предупреждениями/throttle на основе текущего usage из панели. ## Traffic-Тариф @@ -62,7 +63,7 @@ Legacy поле `subscription_duration_months` остается для совм `TariffTrafficWorker` запускается, только если активен `tariffs.json`. -- Раз в несколько минут проверяет наступление нового календарного месяца. +- Раз в несколько минут синхронизирует `trafficLimitStrategy = MONTH` для period-тарифов, если панель ещё не переключена. - Отправляет/дедуплицирует уровни предупреждений 80/95/100 через `traffic_warnings`. - При 100% удаляет пользователя из squad-ов тарифа и ставит `is_throttled`. - Возвращает пользователя в squad-ы, когда лимит снова больше использованного трафика. diff --git a/tests/test_tariff_worker.py b/tests/test_tariff_worker.py index 879d9df..3885cbe 100644 --- a/tests/test_tariff_worker.py +++ b/tests/test_tariff_worker.py @@ -1,62 +1,78 @@ +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 unittest.mock import AsyncMock -from config.settings import Settings +from bot.services.panel_api_service import PanelApiService +from bot.services.subscription_service import SubscriptionService from bot.services.tariff_worker import TariffTrafficWorker -from bot.utils.date_utils import month_start +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_monthly_reset_triggers_on_calendar_change_before_30_days(self): - settings = Settings( - _env_file=None, - BOT_TOKEN="token", - POSTGRES_USER="app_user", - POSTGRES_PASSWORD="app_password", - ) - panel_service = AsyncMock() - panel_service.reset_user_traffic = AsyncMock(return_value=True) - panel_service.add_users_to_internal_squad = AsyncMock(return_value=True) - subscription_service = SimpleNamespace() - worker = TariffTrafficWorker( - settings=settings, - session_factory=SimpleNamespace(), - panel_service=panel_service, - subscription_service=subscription_service, - ) + 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") - now = datetime.now(timezone.utc) - current_month_start = month_start(now) - previous_month_anchor = current_month_start - timedelta(days=2) - sub = SimpleNamespace( - subscription_id=1, - panel_user_uuid="panel-uuid", - period_start_at=previous_month_anchor, - end_date=now + timedelta(days=10), - traffic_used_bytes=123, - traffic_limit_bytes=456, - topup_balance_bytes=0, - is_throttled=False, - status_from_panel="ACTIVE", - ) - tariff = SimpleNamespace(billing_model="period", squad_uuids=["squad-1"]) + 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, + ) - session = object() + 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") - with patch( - "bot.services.tariff_worker.subscription_dal.update_subscription", - new=AsyncMock(), - ) as update_subscription, patch( - "bot.services.tariff_worker.tariff_dal.clear_period_warnings", - new=AsyncMock(), - ) as clear_period_warnings: - await worker._maybe_reset_period(session, sub, tariff, used=123) + await worker._ensure_period_reset_strategy(sub, tariff, sub.traffic_limit_bytes, "NO_RESET") - panel_service.reset_user_traffic.assert_awaited_once_with("panel-uuid") - update_subscription.assert_awaited_once() - update_payload = update_subscription.await_args.args[2] - self.assertEqual(update_payload["period_start_at"], current_month_start) - self.assertEqual(update_payload["traffic_used_bytes"], 0) - clear_period_warnings.assert_awaited_once_with(session, 1) + 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)