fix: preserve premium topup squad access

This commit is contained in:
3252a8
2026-05-25 12:18:58 +03:00
parent 340afe1b80
commit ba7811621e
6 changed files with 466 additions and 26 deletions
@@ -349,24 +349,23 @@ class TrafficMixin:
},
)
panel_payload = {
"uuid": db_user.panel_user_uuid,
"activeInternalSquads": self._panel_squads_for_tariff(
desired_squads = self._panel_squads_for_tariff(
tariff,
include_premium=not premium_is_limited,
),
}
updated_panel = await self.panel_service.update_user_details_on_panel(
db_user.panel_user_uuid, panel_payload
)
if not updated_panel or updated_panel.get("error"):
panel_updated = await self._sync_panel_squads_if_needed(
db_user.panel_user_uuid,
desired_squads,
user_id=user_id,
source="premium_topup",
)
if not panel_updated:
# Otherwise the user pays for premium top-up but the panel never
# re-grants premium squad access (the most common case here is
# transitioning from premium_is_limited=True back to False).
logging.warning(
"Panel user details update FAILED for premium top-up user %s. Response: %s",
"Panel user details update FAILED for premium top-up user %s.",
user_id,
updated_panel,
)
return None
await tariff_dal.create_traffic_topup(
@@ -432,10 +431,16 @@ class TrafficMixin:
squads = self._panel_squads_for_tariff(tariff, include_premium=not premium_is_limited)
try:
await self.panel_service.update_user_details_on_panel(
panel_updated = await self._sync_panel_squads_if_needed(
db_user.panel_user_uuid,
{"uuid": db_user.panel_user_uuid, "activeInternalSquads": squads},
log_response=False,
squads,
user_id=user_id,
source="admin_premium_override",
)
if not panel_updated:
logging.warning(
"sync_premium_squad_access_to_panel: panel update failed for user %s",
user_id,
)
except Exception:
logging.exception(
@@ -679,16 +684,21 @@ class TrafficMixin:
"premium_period_start_at": premium_period_start,
},
)
panel_payload = {
"uuid": db_user.panel_user_uuid,
"activeInternalSquads": self._panel_squads_for_tariff(
desired_squads = self._panel_squads_for_tariff(
tariff,
include_premium=not premium_is_limited,
),
}
)
try:
await self.panel_service.update_user_details_on_panel(
db_user.panel_user_uuid, panel_payload
panel_updated = await self._sync_panel_squads_if_needed(
db_user.panel_user_uuid,
desired_squads,
user_id=user_id,
source="admin_premium_topup",
)
if not panel_updated:
logging.warning(
"admin_grant_premium_topup: panel update failed for user %s",
user_id,
)
except Exception:
logging.exception(
@@ -710,3 +720,128 @@ class TrafficMixin:
"premium_is_limited": premium_is_limited,
"granted_bytes": purchase_bytes,
}
async def _sync_panel_squads_if_needed(
self,
panel_user_uuid: str,
desired_squads: List[str],
*,
user_id: int,
source: str,
) -> bool:
match, current_set = await self._panel_squads_match(panel_user_uuid, desired_squads)
if match is True:
return True
desired_set = self._panel_squad_uuid_set(desired_squads)
self._log_panel_squad_patch(
source=source,
user_id=user_id,
panel_uuid=panel_user_uuid,
current_set=current_set,
desired_set=desired_set,
)
updated_panel = await self.panel_service.update_user_details_on_panel(
panel_user_uuid,
{"uuid": panel_user_uuid, "activeInternalSquads": desired_squads},
log_response=False,
)
if not updated_panel:
return False
return not (isinstance(updated_panel, dict) and updated_panel.get("error"))
async def _panel_squads_match(
self,
panel_user_uuid: str,
desired_squads: List[str],
) -> tuple[Optional[bool], Optional[set[str]]]:
try:
panel_user = await self.panel_service.get_user_by_uuid(
panel_user_uuid,
log_response=False,
)
except Exception:
logging.exception(
"Failed to fetch panel user %s before premium squad update",
panel_user_uuid,
)
return None, None
current_known, current_set = self._panel_active_squad_uuid_set(panel_user)
if not current_known:
return None, current_set
return current_set == self._panel_squad_uuid_set(desired_squads), current_set
@classmethod
def _panel_active_squad_uuid_set(
cls,
panel_user: Optional[dict],
) -> tuple[bool, set[str]]:
if not isinstance(panel_user, dict):
return False, set()
for key in (
"activeInternalSquads",
"active_internal_squads",
"activeInternalSquadUuids",
"active_internal_squad_uuids",
):
if key in panel_user:
return True, cls._panel_squad_uuid_set(panel_user.get(key))
return False, set()
@staticmethod
def _panel_squad_uuid_set(raw) -> set[str]:
if not isinstance(raw, (list, tuple, set)):
return set()
out: set[str] = set()
for item in raw:
if isinstance(item, dict):
nested_squad = item.get("internalSquad") or item.get("squad")
if not isinstance(nested_squad, dict):
nested_squad = {}
squad_uuid = (
item.get("uuid")
or item.get("internalSquadUuid")
or item.get("squadUuid")
or nested_squad.get("uuid")
)
if squad_uuid:
out.add(str(squad_uuid))
elif item:
out.add(str(item))
return out
def _log_panel_squad_patch(
self,
*,
source: str,
user_id: int,
panel_uuid: str,
current_set: Optional[set[str]],
desired_set: set[str],
) -> None:
logging.info(
"Sync panel PATCH: source=%s user_id=%s telegram_id=%s panel_uuid=%s "
"panel_view=full_fetch reasons=activeInternalSquads_mismatch "
"fields=activeInternalSquads payload_fields=activeInternalSquads changes=%s",
source,
user_id,
user_id,
panel_uuid,
"activeInternalSquads:%s->%s"
% (
self._format_panel_squad_set(current_set),
self._format_panel_squad_set(desired_set),
),
)
@staticmethod
def _format_panel_squad_set(value: Optional[set[str]]) -> str:
if value is None:
return "missing"
values = sorted(str(item) for item in value)
preview = ",".join(values[:4])
suffix = ",..." if len(values) > 4 else ""
text = f"[{len(values)}:{preview}{suffix}]"
if len(text) > 96:
return f"{text[:93]}..."
return text
+77 -1
View File
@@ -486,12 +486,22 @@ class TariffTrafficWorker:
return
premium_period_start = month_start(now)
same_period = bool(getattr(sub, "premium_period_start_at", None) == premium_period_start)
same_period = self._same_premium_period(
getattr(sub, "premium_period_start_at", None),
premium_period_start,
)
premium_baseline = int(tariff.premium_monthly_bytes or 0)
premium_topup_balance = int(sub.premium_topup_balance_bytes or 0)
premium_topup_used = (
int(getattr(sub, "premium_topup_used_bytes", 0) or 0) if same_period else 0
)
premium_topup_balance = await self._repair_premium_topup_balance_from_ledger(
session,
sub,
premium_period_start,
premium_topup_balance,
premium_topup_used,
)
# Admin-side overrides for free gifted premium traffic.
premium_unlimited_override = bool(getattr(sub, "premium_unlimited_override", False))
premium_bonus = max(0, int(getattr(sub, "premium_bonus_bytes", 0) or 0))
@@ -648,6 +658,72 @@ class TariffTrafficWorker:
)
return None
@staticmethod
def _same_premium_period(value: Optional[datetime], premium_period_start: datetime) -> bool:
if value is None:
return False
try:
return month_start(value) == premium_period_start
except Exception:
return False
async def _repair_premium_topup_balance_from_ledger(
self,
session: AsyncSession,
sub: Subscription,
premium_period_start: datetime,
premium_topup_balance: int,
premium_topup_used: int,
) -> int:
ledger_total = await self._premium_topup_ledger_total(
session,
int(getattr(sub, "subscription_id", 0) or 0),
premium_period_start,
)
if ledger_total is None:
return premium_topup_balance
tracked_total = max(0, int(premium_topup_balance or 0)) + max(
0,
int(premium_topup_used or 0),
)
if ledger_total <= tracked_total:
return premium_topup_balance
repaired_bytes = ledger_total - tracked_total
logging.warning(
"Premium top-up balance repaired from ledger for user %s subscription %s: "
"tracked=%s ledger=%s repaired=%s",
getattr(sub, "user_id", None),
getattr(sub, "subscription_id", None),
tracked_total,
ledger_total,
repaired_bytes,
)
return premium_topup_balance + repaired_bytes
async def _premium_topup_ledger_total(
self,
session: AsyncSession,
subscription_id: int,
premium_period_start: datetime,
) -> Optional[int]:
if not subscription_id or not isinstance(session, AsyncSession):
return None
try:
return await tariff_dal.sum_traffic_topups(
session,
subscription_id=subscription_id,
kinds=["premium_topup", "admin_premium_topup"],
created_at_gte=premium_period_start,
)
except Exception:
logging.exception(
"TariffTrafficWorker: failed to read premium top-up ledger for subscription %s",
subscription_id,
)
return None
@staticmethod
def _premium_squad_match_cache_key(
panel_user_uuid: str,
+21 -1
View File
@@ -1,6 +1,6 @@
from typing import Any, Dict, List, Optional
from sqlalchemy import and_, delete, select
from sqlalchemy import and_, delete, func, select
from sqlalchemy.ext.asyncio import AsyncSession
from db.models import HwidDevicePurchase, TariffChange, TrafficTopup, TrafficWarning
@@ -26,6 +26,26 @@ async def create_traffic_topup(
return record
async def sum_traffic_topups(
session: AsyncSession,
*,
subscription_id: int,
kinds: Optional[List[str]] = None,
created_at_gte=None,
) -> int:
conditions = [TrafficTopup.subscription_id == subscription_id]
if kinds:
conditions.append(TrafficTopup.kind.in_(list(kinds)))
if created_at_gte is not None:
conditions.append(TrafficTopup.created_at >= created_at_gte)
result = await session.execute(
select(func.coalesce(func.sum(TrafficTopup.purchased_bytes), 0)).where(
and_(*conditions)
)
)
return int(result.scalar() or 0)
async def create_hwid_device_purchase(
session: AsyncSession,
*,
+73
View File
@@ -209,6 +209,79 @@ class AdminGrantTopupTests(unittest.IsolatedAsyncioTestCase):
sub_update_payload = upd.await_args.args[2]
self.assertFalse(sub_update_payload["premium_is_limited"])
async def test_premium_grant_skips_panel_patch_when_squads_already_match(self):
with tempfile.TemporaryDirectory() as tmpdir:
settings = _make_settings(_tariffs_config_payload(premium=True), tmpdir)
panel_service = AsyncMock(spec=PanelApiService)
panel_service.get_user_by_uuid = AsyncMock(
return_value={
"activeInternalSquads": [
{"uuid": "squad-1"},
{"uuid": "premium-squad"},
]
}
)
panel_service.update_user_details_on_panel = AsyncMock(return_value={"response": {}})
service = SubscriptionService(settings, panel_service)
db_user = SimpleNamespace(
user_id=77,
first_name="Premium",
last_name=None,
username="premium",
language_code="ru",
panel_user_uuid="panel-uuid",
email=None,
telegram_id=77,
)
sub = SimpleNamespace(
subscription_id=9,
user_id=77,
panel_user_uuid="panel-uuid",
tariff_key="standard",
premium_baseline_bytes=25 * (1024**3),
premium_topup_balance_bytes=0,
premium_topup_used_bytes=0,
premium_used_bytes=30 * (1024**3),
premium_is_limited=True,
premium_period_start_at=datetime.now(timezone.utc).replace(
day=1, hour=0, minute=0, second=0, microsecond=0
),
premium_unlimited_override=False,
premium_bonus_bytes=0,
)
updated_sub = SimpleNamespace(**vars(sub))
with (
patch(
"bot.services.subscription_service.user_dal.get_user_by_id",
new=AsyncMock(return_value=db_user),
),
patch(
"bot.services.subscription_service.subscription_dal.get_active_subscription_by_user_id",
new=AsyncMock(return_value=sub),
),
patch(
"bot.services.subscription_service.subscription_dal.update_subscription",
new=AsyncMock(return_value=updated_sub),
) as upd,
patch(
"bot.services.subscription_service.tariff_dal.create_traffic_topup",
new=AsyncMock(),
) as topup_log,
):
result = await service.admin_grant_premium_topup(AsyncMock(), 77, 20.0)
self.assertIsNotNone(result)
self.assertFalse(result["premium_is_limited"])
upd.assert_awaited_once()
topup_log.assert_awaited_once()
panel_service.get_user_by_uuid.assert_awaited_once_with(
"panel-uuid",
log_response=False,
)
panel_service.update_user_details_on_panel.assert_not_awaited()
async def test_premium_grant_fails_when_tariff_has_no_premium_squads(self):
with tempfile.TemporaryDirectory() as tmpdir:
settings = _make_settings(_tariffs_config_payload(premium=False), tmpdir)
+75
View File
@@ -277,6 +277,81 @@ class TariffWorkerTests(unittest.IsolatedAsyncioTestCase):
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_topup_ledger_repairs_missing_balance_before_limiting(self):
payload = _tariffs_config_payload()
payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"]
payload["tariffs"][0]["premium_monthly_gb"] = 25
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": 40 * (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,
)
worker._premium_topup_ledger_total = AsyncMock(return_value=20 * (1024**3))
now = datetime(2026, 5, 9, tzinfo=timezone.utc)
sub = SimpleNamespace(
subscription_id=1,
user_id=123,
panel_user_uuid="panel-uuid",
premium_baseline_bytes=25 * (1024**3),
premium_topup_balance_bytes=0,
premium_topup_used_bytes=0,
premium_used_bytes=40 * (1024**3),
premium_is_limited=True,
premium_period_start_at=datetime(2026, 5, 1, tzinfo=timezone.utc),
premium_unlimited_override=False,
premium_bonus_bytes=0,
)
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,
now,
panel_username="tg_123",
panel_user_dict={
"activeInternalSquads": [
{"uuid": "squad-1"},
{"uuid": "premium-squad"},
]
},
)
self.assertEqual(sub.premium_topup_balance_bytes, 5 * (1024**3))
self.assertEqual(sub.premium_topup_used_bytes, 15 * (1024**3))
self.assertFalse(sub.premium_is_limited)
panel_service.update_user_details_on_panel.assert_not_awaited()
async def test_premium_usage_update_does_not_patch_panel_when_access_state_unchanged(self):
payload = _tariffs_config_payload()
payload["tariffs"][0]["premium_squad_uuids"] = ["premium-squad"]
+61
View File
@@ -270,6 +270,67 @@ class ActivatePremiumTopupPanelFailureTests(unittest.IsolatedAsyncioTestCase):
self.assertEqual(result["tariff_key"], "standard")
create_topup.assert_awaited_once()
async def test_skips_panel_patch_when_premium_squads_already_match(self):
with tempfile.TemporaryDirectory() as tmpdir:
settings = _make_settings(tmpdir)
service = _make_service(settings)
sub = _make_sub(
premium_topup_balance_bytes=0,
premium_used_bytes=30 * GIB,
premium_is_limited=True,
)
user = _make_user()
service.panel_service.get_user_by_uuid = AsyncMock(
return_value={
"activeInternalSquads": [
{"uuid": "main-squad"},
{"uuid": "premium-squad"},
]
}
)
service.panel_service.update_user_details_on_panel = AsyncMock(
return_value={"ok": True, "uuid": "panel-uuid"}
)
with (
patch(
"bot.services.subscription_service_impl.payments.payment_dal.get_payment_by_db_id",
AsyncMock(return_value=SimpleNamespace()),
),
patch(
"bot.services.subscription_service_impl.traffic.user_dal.get_user_by_id",
AsyncMock(return_value=user),
),
patch(
"bot.services.subscription_service_impl.traffic.subscription_dal.get_active_subscription_by_user_id",
AsyncMock(return_value=sub),
),
patch(
"bot.services.subscription_service_impl.traffic.subscription_dal.update_subscription",
AsyncMock(),
),
patch(
"bot.services.subscription_service_impl.traffic.tariff_dal.create_traffic_topup",
AsyncMock(),
) as create_topup,
):
result = await service.activate_premium_topup(
session=AsyncMock(),
user_id=42,
tariff_key="standard",
traffic_gb=10,
payment_amount=100,
payment_db_id=2,
)
self.assertIsNotNone(result)
create_topup.assert_awaited_once()
service.panel_service.get_user_by_uuid.assert_awaited_once_with(
"panel-uuid",
log_response=False,
)
service.panel_service.update_user_details_on_panel.assert_not_awaited()
class SwitchTariffPanelFailureTests(unittest.IsolatedAsyncioTestCase):
"""Free tariff switch. The bug was lifecycle.py:121 ignoring panel result — the tariff_key