server/benefit: add a flag on benefit service to check if we should r…

…evoke benefit individually or not

Useful for license keys where we generate a single key for each grant, contrary to other benefits where the side effect is global and should only revoked for the last remaining grant (like GitHub Repository invitation)

server/polar/benefit/benefits/base.py

@@ -84,6 +84,8 @@ class BenefitServiceProtocol(Protocol[B, BP, BGP]):
     session: AsyncSession
     redis: Redis
 
+    should_revoke_individually: bool = False
+
     def __init__(self, session: AsyncSession, redis: Redis) -> None:
         self.session = session
         self.redis = redis

server/polar/benefit/benefits/license_keys.py

@@ -26,6 +26,8 @@ class BenefitLicenseKeysService(
         BenefitGrantLicenseKeysProperties,
     ]
 ):
+    should_revoke_individually = True
+
     async def grant(
         self,
         benefit: BenefitLicenseKeys,

server/polar/benefit/service/benefit_grant.py

@@ -230,13 +230,15 @@ async def revoke_benefit(
 
         previous_properties = grant.properties
 
-        # Check if the user has other grants (from other purchases) for this benefit
-        # If yes, don't call the revoke logic, just mark the grant as revoked
+        benefit_service = get_benefit_service(benefit.type, session, redis)
+        # Call the revoke logic in two cases:
+        # * If the service requires grants to be revoked individually
+        # * If there is only one grant remaining for this benefit,
+        #   so the benefit remains if other grants exist via other purchases
         other_grants = await self._get_granted_by_benefit_and_user(
             session, benefit, user
         )
-        if len(other_grants) < 2:
-            benefit_service = get_benefit_service(benefit.type, session, redis)
+        if benefit_service.should_revoke_individually or len(other_grants) < 2:
             properties = await benefit_service.revoke(
                 benefit,
                 user,

server/tests/benefit/service/test_benefit_grant.py

@@ -31,6 +31,7 @@
 @pytest.fixture(autouse=True)
 def benefit_service_mock(mocker: MockerFixture) -> MagicMock:
     service_mock = MagicMock(spec=BenefitServiceProtocol)
+    service_mock.should_revoke_individually = False
     service_mock.grant.return_value = {}
     service_mock.revoke.return_value = {}
     mock = mocker.patch("polar.benefit.service.benefit_grant.get_benefit_service")
@@ -286,6 +287,47 @@ async def test_several_benefit_grants(
         assert updated_grant.is_revoked
         benefit_service_mock.revoke.assert_not_called()
 
+    async def test_several_benefit_grants_should_individual_revoke(
+        self,
+        session: AsyncSession,
+        redis: Redis,
+        save_fixture: SaveFixture,
+        subscription: Subscription,
+        user: User,
+        benefit_organization: Benefit,
+        benefit_service_mock: MagicMock,
+        product: Product,
+    ) -> None:
+        benefit_service_mock.should_revoke_individually = True
+        benefit_service_mock.revoke.return_value = {"message": "ok"}
+
+        first_grant = await create_benefit_grant(
+            save_fixture, user, benefit_organization, subscription=subscription
+        )
+        first_grant.set_granted()
+        await save_fixture(first_grant)
+
+        second_subscription = await create_subscription(
+            save_fixture, product=product, user=user
+        )
+        second_grant = await create_benefit_grant(
+            save_fixture, user, benefit_organization, subscription=second_subscription
+        )
+        second_grant.set_granted()
+        await save_fixture(second_grant)
+
+        # then
+        session.expunge_all()
+
+        updated_grant = await benefit_grant_service.revoke_benefit(
+            session, redis, user, benefit_organization, subscription=subscription
+        )
+
+        assert updated_grant.id == first_grant.id
+        assert updated_grant.is_revoked
+        assert updated_grant.properties == {"message": "ok"}
+        benefit_service_mock.revoke.assert_called_once()
+
 
 @pytest.mark.asyncio
 @pytest.mark.skip_db_asserts