feat(auto-redeem): Support ZZZ

seriaati · Jul 5, 2024 · eb4bc0c · eb4bc0c
1 parent 35a8d6a
commit eb4bc0c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 18 deletions.
diff --git a/hoyo_buddy/cogs/admin.py b/hoyo_buddy/cogs/admin.py
@@ -8,7 +8,7 @@
 from genshin import Game  # noqa: TCH002
 from seria.utils import write_json
 
-from ..constants import GPY_GAME_TO_HB_GAME, UID_STARTS
+from ..constants import UID_STARTS
 from ..hoyo.auto_tasks.auto_redeem import AutoRedeem
 from ..hoyo.auto_tasks.daily_checkin import DailyCheckin
 from ..hoyo.auto_tasks.farm_check import FarmChecker
@@ -140,7 +140,7 @@ async def add_codes_command(self, ctx: commands.Context, game: Game, codes: str)
             codes_.add(code_)
 
         message = await ctx.send("Adding codes...")
-        task_ran = await AutoRedeem.execute(self.bot, GPY_GAME_TO_HB_GAME[game], codes.split(","))
+        task_ran = await AutoRedeem.execute(self.bot, game, codes.split(","))
         if not task_ran:
             await message.edit(content="Auto redeem task is already running.")
 

diff --git a/hoyo_buddy/hoyo/auto_tasks/auto_redeem.py b/hoyo_buddy/hoyo/auto_tasks/auto_redeem.py
@@ -9,8 +9,9 @@
 
 from ...bot.error_handler import get_error_embed
 from ...bot.translator import LocaleStr
+from ...constants import HB_GAME_TO_GPY_GAME
 from ...db.models import HoyoAccount
-from ...enums import Game, Platform
+from ...enums import Platform
 
 if TYPE_CHECKING:
     import aiohttp
@@ -85,7 +86,7 @@ async def _handle_error(cls, account: HoyoAccount, locale: Locale, e: Exception)
 
     @classmethod
     async def execute(
-        cls, bot: HoyoBuddy, game: Game | None = None, codes: list[str] | None = None
+        cls, bot: HoyoBuddy, game: genshin.Game | None = None, codes: list[str] | None = None
     ) -> bool:
         """Redeem codes for accounts that have auto redeem enabled.
 
@@ -106,12 +107,11 @@ async def execute(
             )
             cls._bot = bot
 
-            if codes is None:
-                genshin_codes = await cls._get_codes(bot.session, genshin.Game.GENSHIN)
-                hsr_codes = await cls._get_codes(bot.session, genshin.Game.STARRAIL)
-            else:
-                genshin_codes = codes
-                hsr_codes = codes
+            game_codes = (
+                {game: codes}
+                if game is not None and codes is not None
+                else {game_: await cls._get_codes(bot.session, game_) for game_ in genshin.Game}
+            )
 
             accounts = (
                 await HoyoAccount.filter(auto_redeem=True, game=game).all()
@@ -123,14 +123,9 @@ async def execute(
                 if account.platform is Platform.MIYOUSHE:
                     continue
 
-                if account.game is Game.GENSHIN:
-                    codes = genshin_codes
-                elif account.game is Game.STARRAIL:
-                    codes = hsr_codes
-                else:
-                    continue
-
-                await cls._redeem_codes(codes, account)
+                await cls._redeem_codes(
+                    game_codes.get(HB_GAME_TO_GPY_GAME[account.game], []), account
+                )
 
             logger.info("Auto redeem task completed")