chore: cover referral bonus skip-paths, inviter and referee award flows

tests/test_referral_bonuses.py

@@ -0,0 +1,427 @@
+"""Behavioural tests for ``ReferralService.apply_referral_bonuses_for_payment``.
+
+This is the entry point that every payment provider calls after a successful
+``subscription`` purchase. The contract has several non-obvious branches that
+have been wrong before in this project:
+
+* skip when the referee has no inviter recorded;
+* skip when ``REFERRAL_ONE_BONUS_PER_REFEREE`` is set and the referee already
+  has prior succeeded payments;
+* skip when the referee was already active at payment time and the caller
+  asked to skip that case;
+* award both the inviter bonus (extend existing or create new bonus sub)
+  and the referee bonus (extend their freshly activated subscription).
+
+The bot/i18n/panel collaborators are stubbed because we are testing branching
+and DB intent, not network I/O.
+"""
+
+import unittest
+from datetime import datetime, timezone
+from types import SimpleNamespace
+from typing import Any, Optional
+from unittest.mock import AsyncMock, patch
+
+from bot.services.referral_service import ReferralService
+
+
+def _make_settings(**overrides: Any) -> SimpleNamespace:
+    base = dict(
+        DEFAULT_LANGUAGE="en",
+        REFERRAL_ONE_BONUS_PER_REFEREE=True,
+        user_traffic_limit_bytes=0,
+        referral_bonus_inviter={1: 7, 3: 21, 6: 45, 12: 90},
+        referral_bonus_referee={1: 3, 3: 10, 6: 21, 12: 45},
+    )
+    base.update(overrides)
+    return SimpleNamespace(**base)
+
+
+def _make_user(user_id: int, *, referred_by_id: Optional[int] = None, **fields):
+    base = dict(
+        user_id=user_id,
+        first_name=f"User{user_id}",
+        language_code="en",
+        referred_by_id=referred_by_id,
+        email=None,
+        panel_user_uuid=f"panel-{user_id}",
+    )
+    base.update(fields)
+    return SimpleNamespace(**base)
+
+
+def _make_service(*, settings, subscription_service):
+    bot = AsyncMock()
+    bot.send_message = AsyncMock()
+    i18n = SimpleNamespace(gettext=lambda lang, key, **kw: f"{key}({lang})")
+    return ReferralService(settings, subscription_service, bot, i18n), bot
+
+
+class SkipPathTests(unittest.IsolatedAsyncioTestCase):
+    async def test_no_inviter_returns_empty_payload(self):
+        settings = _make_settings()
+        subscription_service = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+        # Referred_by_id is None → bail out immediately.
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(return_value=_make_user(42, referred_by_id=None)),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+            )
+        self.assertEqual(
+            result,
+            {"referee_bonus_applied_days": None, "referee_new_end_date": None},
+        )
+        subscription_service.extend_active_subscription_days.assert_not_called()
+
+    async def test_referee_with_prior_successful_payment_is_skipped(self):
+        # REFERRAL_ONE_BONUS_PER_REFEREE protects against duplicate awarding
+        # when the same referee buys multiple subscriptions.
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=True)
+        subscription_service = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+        with (
+            patch(
+                "bot.services.referral_service.user_dal.get_user_by_id",
+                AsyncMock(
+                    side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                    if uid == 42
+                    else _make_user(uid)
+                ),
+            ),
+            patch(
+                "bot.services.referral_service.payment_dal.count_user_succeeded_payments",
+                AsyncMock(return_value=2),
+            ),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+                current_payment_db_id=99,
+            )
+        self.assertEqual(
+            result,
+            {"referee_bonus_applied_days": None, "referee_new_end_date": None},
+        )
+        subscription_service.extend_active_subscription_days.assert_not_called()
+
+    async def test_referee_already_active_is_skipped_when_caller_asks(self):
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=False)
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=True)
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(
+                side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                if uid == 42
+                else _make_user(uid)
+            ),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+                skip_if_active_before_payment=True,
+            )
+        self.assertEqual(
+            result,
+            {"referee_bonus_applied_days": None, "referee_new_end_date": None},
+        )
+        subscription_service.extend_active_subscription_days.assert_not_called()
+
+
+class InviterBonusTests(unittest.IsolatedAsyncioTestCase):
+    """The inviter bonus has two sub-paths: extend their active sub or create
+    a fresh bonus sub when they had none. Pin both."""
+
+    async def test_extends_inviter_active_subscription(self):
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=False)
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=False)
+        # _get_or_create_panel_user_link_details returns (uuid, sub_link_id, short, created)
+        subscription_service._get_or_create_panel_user_link_details = AsyncMock(
+            return_value=("inviter-panel", "inviter-sub", "short", False)
+        )
+        new_end = datetime(2026, 1, 1, tzinfo=timezone.utc)
+        subscription_service.extend_active_subscription_days = AsyncMock(return_value=new_end)
+        service, bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(
+                side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                if uid == 42
+                else _make_user(uid)
+            ),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+                skip_if_active_before_payment=False,
+            )
+
+        self.assertTrue(result["inviter_bonus_applied_flag"])
+        # Inviter extension call must use the configured 1-month inviter bonus (7 days here).
+        subscription_service.extend_active_subscription_days.assert_any_await(
+            session=unittest.mock.ANY,
+            user_id=1,
+            bonus_days=7,
+            reason=unittest.mock.ANY,
+        )
+        bot.send_message.assert_any_await(1, unittest.mock.ANY)
+
+    async def test_creates_new_bonus_sub_when_inviter_has_no_active_sub(self):
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=False)
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=False)
+        subscription_service._get_or_create_panel_user_link_details = AsyncMock(
+            return_value=("inviter-panel", "inviter-sub", "short", False)
+        )
+        # Inviter has no active sub → extend returns None for inviter.
+        # extend_active_subscription_days is called twice (inviter + referee). Use a side_effect:
+        # First call (inviter) → None, second call (referee) → new datetime.
+        referee_new_end = datetime(2026, 2, 1, tzinfo=timezone.utc)
+        subscription_service.extend_active_subscription_days = AsyncMock(
+            side_effect=[None, referee_new_end]
+        )
+        subscription_service.panel_service = AsyncMock()
+        subscription_service.panel_service.update_user_details_on_panel = AsyncMock(
+            return_value={"ok": True}
+        )
+
+        service, bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with (
+            patch(
+                "bot.services.referral_service.user_dal.get_user_by_id",
+                AsyncMock(
+                    side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                    if uid == 42
+                    else _make_user(uid)
+                ),
+            ),
+            patch(
+                "bot.services.referral_service.subscription_dal.deactivate_other_active_subscriptions",
+                AsyncMock(),
+            ),
+            patch(
+                "bot.services.referral_service.subscription_dal.upsert_subscription",
+                AsyncMock(return_value=SimpleNamespace(subscription_id=999)),
+            ) as upsert,
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+                skip_if_active_before_payment=False,
+            )
+
+        self.assertTrue(result["inviter_bonus_applied_flag"])
+        # A new bonus subscription should have been upserted for the inviter.
+        upsert.assert_awaited_once()
+        payload = upsert.await_args.args[1]
+        self.assertEqual(payload["user_id"], 1)
+        self.assertEqual(payload["panel_user_uuid"], "inviter-panel")
+        self.assertEqual(payload["duration_months"], 0)
+        self.assertEqual(payload["status_from_panel"], "ACTIVE_BONUS")
+        # The panel must have been told the new expiry & ACTIVE status.
+        panel_update = subscription_service.panel_service.update_user_details_on_panel
+        panel_update.assert_awaited_once()
+        pn_payload = panel_update.await_args.args[1]
+        self.assertEqual(pn_payload["status"], "ACTIVE")
+        self.assertTrue(pn_payload["expireAt"].endswith("Z"))
+
+
+class RefereeBonusTests(unittest.IsolatedAsyncioTestCase):
+    async def test_applies_referee_bonus_via_extend_active_subscription(self):
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=False)
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=False)
+        # Inviter bonus is 0 for this test (skip inviter branch), referee gets 10 days.
+        subscription_service._get_or_create_panel_user_link_details = AsyncMock(
+            return_value=("inviter-panel", "inviter-sub", "short", False)
+        )
+        referee_new_end = datetime(2026, 3, 1, tzinfo=timezone.utc)
+        subscription_service.extend_active_subscription_days = AsyncMock(
+            return_value=referee_new_end
+        )
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(
+                side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                if uid == 42
+                else _make_user(uid)
+            ),
+        ):
+            # 3-month plan → 10-day referee bonus, 21-day inviter bonus.
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=3,
+                skip_if_active_before_payment=False,
+            )
+
+        self.assertEqual(result["referee_bonus_applied_days"], 10)
+        self.assertEqual(result["referee_new_end_date"], referee_new_end)
+        # Referee extension call:
+        referee_call = [
+            call
+            for call in subscription_service.extend_active_subscription_days.await_args_list
+            if call.kwargs.get("user_id") == 42
+        ]
+        self.assertTrue(referee_call, "referee bonus extension must be invoked")
+        self.assertEqual(referee_call[0].kwargs["bonus_days"], 10)
+
+    async def test_no_referee_bonus_when_extend_returns_none(self):
+        settings = _make_settings(
+            REFERRAL_ONE_BONUS_PER_REFEREE=False,
+            # Skip the inviter branch entirely.
+            referral_bonus_inviter={},
+            referral_bonus_referee={1: 3},
+        )
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=False)
+        subscription_service.extend_active_subscription_days = AsyncMock(return_value=None)
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(
+                side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                if uid == 42
+                else _make_user(uid)
+            ),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=1,
+                skip_if_active_before_payment=False,
+            )
+
+        self.assertIsNone(result["referee_bonus_applied_days"])
+        self.assertIsNone(result["referee_new_end_date"])
+
+    async def test_no_bonus_for_months_without_configured_value(self):
+        # 24-month plan has no entries in either bonus dict → nothing applies.
+        settings = _make_settings(REFERRAL_ONE_BONUS_PER_REFEREE=False)
+        subscription_service = AsyncMock()
+        subscription_service.has_active_subscription = AsyncMock(return_value=False)
+        subscription_service.extend_active_subscription_days = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(
+                side_effect=lambda session, uid: _make_user(uid, referred_by_id=1)
+                if uid == 42
+                else _make_user(uid)
+            ),
+        ):
+            result = await service.apply_referral_bonuses_for_payment(
+                session=AsyncMock(),
+                referee_user_id=42,
+                purchased_subscription_months=24,
+                skip_if_active_before_payment=False,
+            )
+
+        self.assertIsNone(result["referee_bonus_applied_days"])
+        self.assertFalse(result["inviter_bonus_applied_flag"])
+        subscription_service.extend_active_subscription_days.assert_not_called()
+
+
+class GenerateReferralLinkTests(unittest.IsolatedAsyncioTestCase):
+    async def test_includes_bot_username_and_referral_code(self):
+        settings = _make_settings()
+        subscription_service = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with (
+            patch(
+                "bot.services.referral_service.user_dal.get_user_by_id",
+                AsyncMock(return_value=_make_user(1)),
+            ),
+            patch(
+                "bot.services.referral_service.user_dal.ensure_referral_code",
+                AsyncMock(return_value="abc123"),
+            ),
+        ):
+            link = await service.generate_referral_link(
+                session=AsyncMock(),
+                bot_username="my_bot",
+                inviter_user_id=1,
+            )
+        self.assertEqual(link, "https://t.me/my_bot?start=ref_uabc123")
+
+    async def test_returns_none_when_user_missing(self):
+        settings = _make_settings()
+        subscription_service = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with patch(
+            "bot.services.referral_service.user_dal.get_user_by_id",
+            AsyncMock(return_value=None),
+        ):
+            link = await service.generate_referral_link(
+                session=AsyncMock(),
+                bot_username="my_bot",
+                inviter_user_id=999,
+            )
+        self.assertIsNone(link)
+
+    async def test_returns_none_when_referral_code_unavailable(self):
+        settings = _make_settings()
+        subscription_service = AsyncMock()
+        service, _bot = _make_service(
+            settings=settings, subscription_service=subscription_service
+        )
+
+        with (
+            patch(
+                "bot.services.referral_service.user_dal.get_user_by_id",
+                AsyncMock(return_value=_make_user(1)),
+            ),
+            patch(
+                "bot.services.referral_service.user_dal.ensure_referral_code",
+                AsyncMock(return_value=None),
+            ),
+        ):
+            link = await service.generate_referral_link(
+                session=AsyncMock(),
+                bot_username="my_bot",
+                inviter_user_id=1,
+            )
+        self.assertIsNone(link)
+
+
+if __name__ == "__main__":  # pragma: no cover
+    unittest.main()