From 65144bb46945aa689839e17f472a4382ee227298 Mon Sep 17 00:00:00 2001
From: omg-xtao <100690902+omg-xtao@users.noreply.github.com>
Date: Fri, 27 Jan 2023 11:50:47 +0800
Subject: [PATCH] fix: CookiesNotFoundError when setcookie

* fix: CookiesNotFoundError when setcookie

* remove cookies_database_data
---
 core/cookies/services.py   |  7 +++++++
 plugins/genshin/cookies.py | 30 +++++++++---------------------
 2 files changed, 16 insertions(+), 21 deletions(-)

diff --git a/core/cookies/services.py b/core/cookies/services.py
index 4769f185..d5f54c22 100644
--- a/core/cookies/services.py
+++ b/core/cookies/services.py
@@ -27,6 +27,13 @@ class CookiesService:
     async def del_cookies(self, user_id: int, region: RegionEnum):
         return await self._repository.del_cookies(user_id, region)
 
+    async def add_or_update_cookies(self, user_id: int, cookies: dict, region: RegionEnum):
+        try:
+            await self.get_cookies(user_id, region)
+            await self.update_cookies(user_id, cookies, region)
+        except CookiesNotFoundError:
+            await self.add_cookies(user_id, cookies, region)
+
 
 class PublicCookiesService:
     def __init__(self, cookies_repository: CookiesRepository, public_cookies_cache: PublicCookiesCache):
diff --git a/plugins/genshin/cookies.py b/plugins/genshin/cookies.py
index 45ba180a..c10003c0 100644
--- a/plugins/genshin/cookies.py
+++ b/plugins/genshin/cookies.py
@@ -27,7 +27,6 @@ from utils.models.base import RegionEnum
 
 class AddUserCommandData(TelegramObject):
     user: Optional[User] = None
-    cookies_database_data: Optional[Cookies] = None
     region: RegionEnum = RegionEnum.NULL
     cookies: dict = {}
     game_uid: int = 0
@@ -75,10 +74,8 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
         user = update.effective_user
         message = update.effective_message
         logger.info("用户 %s[%s] 绑定账号命令请求", user.full_name, user.id)
-        add_user_command_data: AddUserCommandData = context.chat_data.get("add_user_command_data")
-        if add_user_command_data is None:
-            cookies_command_data = AddUserCommandData()
-            context.chat_data["add_user_command_data"] = cookies_command_data
+        cookies_command_data = AddUserCommandData()
+        context.chat_data["add_user_command_data"] = cookies_command_data
 
         text = f'你好 {user.mention_markdown_v2()} {escape_markdown("！请选择要绑定的服务器！或回复退出取消操作")}'
         reply_keyboard = [["米游社", "HoYoLab"], ["退出"]]
@@ -92,10 +89,8 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
         user = update.effective_user
         message = update.effective_message
         logger.info("用户 %s[%s] 绑定账号命令请求", user.full_name, user.id)
-        add_user_command_data: AddUserCommandData = context.chat_data.get("add_user_command_data")
-        if add_user_command_data is None:
-            add_user_command_data = AddUserCommandData()
-            context.chat_data["add_user_command_data"] = add_user_command_data
+        add_user_command_data = AddUserCommandData()
+        context.chat_data["add_user_command_data"] = add_user_command_data
         add_user_command_data.region = RegionEnum.HYPERION
         try:
             user_info = await self.user_service.get_user_by_id(user.id)
@@ -103,8 +98,7 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
             user_info = None
         if user_info is not None:
             try:
-                cookies_database_data = await self.cookies_service.get_cookies(user.id, RegionEnum.HYPERION)
-                add_user_command_data.cookies_database_data = cookies_database_data
+                await self.cookies_service.get_cookies(user.id, RegionEnum.HYPERION)
             except CookiesNotFoundError:
                 await message.reply_text("你已经绑定UID，如果继续操作会覆盖当前UID。")
             else:
@@ -149,8 +143,7 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
             user_info = None
         if user_info is not None:
             try:
-                cookies_database_data = await self.cookies_service.get_cookies(user.id, region)
-                add_user_command_data.cookies_database_data = cookies_database_data
+                await self.cookies_service.get_cookies(user.id, region)
             except CookiesNotFoundError:
                 await message.reply_text("你已经绑定UID，如果继续操作会覆盖当前UID。")
             else:
@@ -346,14 +339,9 @@ class SetUserCookies(Plugin.Conversation, BasePlugin.Conversation):
                     await message.reply_text("数据错误")
                     return ConversationHandler.END
                 await self.user_service.update_user(user_db)
-            if add_user_command_data.cookies_database_data is None:
-                await self.cookies_service.add_cookies(
-                    user.id, add_user_command_data.cookies, add_user_command_data.region
-                )
-            else:
-                await self.cookies_service.update_cookies(
-                    user.id, add_user_command_data.cookies, add_user_command_data.region
-                )
+            await self.cookies_service.add_or_update_cookies(
+                user.id, add_user_command_data.cookies, add_user_command_data.region
+            )
             logger.info("用户 %s[%s] 绑定账号成功", user.full_name, user.id)
             await message.reply_text("保存成功", reply_markup=ReplyKeyboardRemove())
             return ConversationHandler.END