Kuby/test_premium.py

import unittest
from unittest.mock import AsyncMock, MagicMock, patch
import os
import json
import asyncio

# Setup env variables before importing premium module
os.environ["STARTER_SKU_ID"] = "111111"
os.environ["MEDIUM_SKU_ID"] = "222222"
os.environ["LARGE_SKU_ID"] = "333333"

from utils.premium import load_partners_whitelist, get_sku_ids, check_premium_tier, WHITELIST_PARTNERS_FILE

class TestPremiumSystem(unittest.IsolatedAsyncioTestCase):
    
    def setUp(self):
        # Reset whitelist file before each test
        if os.path.exists(WHITELIST_PARTNERS_FILE):
            try:
                os.remove(WHITELIST_PARTNERS_FILE)
            except Exception:
                pass

    def tearDown(self):
        # Clean up whitelist file
        if os.path.exists(WHITELIST_PARTNERS_FILE):
            try:
                os.remove(WHITELIST_PARTNERS_FILE)
            except Exception:
                pass

    async def test_get_sku_ids(self):
        starter, medium, large = get_sku_ids()
        self.assertEqual(starter, 111111)
        self.assertEqual(medium, 222222)
        self.assertEqual(large, 333333)

    async def test_load_partners_whitelist_empty(self):
        partners = await load_partners_whitelist()
        self.assertEqual(partners, [])
        self.assertTrue(os.path.exists(WHITELIST_PARTNERS_FILE))

    async def test_load_partners_whitelist_with_data(self):
        with open(WHITELIST_PARTNERS_FILE, "w") as f:
            json.dump({"partners": [123, 456]}, f)
        partners = await load_partners_whitelist()
        self.assertEqual(partners, [123, 456])

    async def test_check_premium_tier_dm_fails(self):
        interaction = MagicMock()
        interaction.guild = None
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertFalse(result)
        interaction.response.send_message.assert_called_once_with(
            "❌ Cette commande ne peut être exécutée que sur un serveur Discord.",
            ephemeral=True
        )

    async def test_check_premium_tier_whitelist_passes(self):
        with open(WHITELIST_PARTNERS_FILE, "w") as f:
            json.dump({"partners": [12345]}, f)
            
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 12345
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertTrue(result)
        interaction.response.send_message.assert_not_called()

    async def test_check_premium_tier_invalid_member_count_fails(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = None
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertFalse(result)
        interaction.response.send_message.assert_called_once_with(
            "❌ Impossible de vérifier la taille du serveur par sécurité. Veuillez réessayer plus tard.",
            ephemeral=True
        )

    async def test_check_premium_tier_less_than_500_with_starter(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = 100
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        # Mock entitlements
        entitlement = MagicMock()
        entitlement.sku_id = 111111
        entitlement.is_active.return_value = True
        interaction.entitlements = [entitlement]
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertTrue(result)

    async def test_check_premium_tier_500_to_2500_with_starter_only_fails(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = 1000
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        entitlement = MagicMock()
        entitlement.sku_id = 111111
        entitlement.is_active.return_value = True
        interaction.entitlements = [entitlement]
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertFalse(result)
        interaction.response.send_message.assert_called_once()
        self.assertIn("Palier Medium", interaction.response.send_message.call_args[0][0])

    async def test_check_premium_tier_500_to_2500_with_medium_passes(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = 1000
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        entitlement = MagicMock()
        entitlement.sku_id = 222222
        entitlement.is_active.return_value = True
        interaction.entitlements = [entitlement]
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertTrue(result)

    async def test_check_premium_tier_above_2500_with_large_passes(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = 3000
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        entitlement = MagicMock()
        entitlement.sku_id = 333333
        entitlement.is_active.return_value = True
        interaction.entitlements = [entitlement]
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertTrue(result)

    async def test_check_premium_tier_above_2500_with_medium_fails(self):
        interaction = MagicMock()
        interaction.guild = MagicMock()
        interaction.guild.id = 999
        interaction.guild.member_count = 3000
        interaction.response = MagicMock()
        interaction.response.send_message = AsyncMock()
        
        entitlement = MagicMock()
        entitlement.sku_id = 222222
        entitlement.is_active.return_value = True
        interaction.entitlements = [entitlement]
        
        decorator = check_premium_tier()
        predicate = decorator.predicate
        
        result = await predicate(interaction)
        self.assertFalse(result)
        self.assertIn("Palier Large", interaction.response.send_message.call_args[0][0])

if __name__ == "__main__":
    unittest.main()