From fb6c3a9f36b8a5d112c9342e75250c7e09dc66e5 Mon Sep 17 00:00:00 2001
From: jxxghp <jxxghp@gmail.com>
Date: Wed, 27 Mar 2024 15:45:27 +0800
Subject: [PATCH] fix site test

---
 app/api/endpoints/site.py |  2 +-
 app/chain/site.py         | 40 +++++++++++++++++++++------------------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/app/api/endpoints/site.py b/app/api/endpoints/site.py
index 9adb38ea..c5623e77 100644
--- a/app/api/endpoints/site.py
+++ b/app/api/endpoints/site.py
@@ -174,7 +174,7 @@ def test_site(site_id: int,
             status_code=404,
             detail=f"站点 {site_id} 不存在",
         )
-    status, message = SiteChain().test(url=site.domain)
+    status, message = SiteChain().test(site.domain)
     return schemas.Response(success=status, message=message)
 
 
diff --git a/app/chain/site.py b/app/chain/site.py
index a0490a13..f5db1df5 100644
--- a/app/chain/site.py
+++ b/app/chain/site.py
@@ -185,7 +185,7 @@ class SiteChain(ChainBase):
             site_info = self.siteoper.get_by_domain(domain)
             if site_info:
                 # 站点已存在，检查站点连通性
-                status, msg = self.test(url=domain)
+                status, msg = self.test(domain)
                 # 更新站点Cookie
                 if status:
                     logger.info(f"站点【{site_info.name}】连通性正常，不同步CookieCloud数据")
@@ -327,29 +327,19 @@ class SiteChain(ChainBase):
                 logger.info(f"清理站点配置：{key}")
                 self.systemconfig.delete(key)
 
-    def test(self, url: str = None, site_info: CommentedMap = None) -> Tuple[bool, str]:
+    def test(self, site_info: Union[str, CommentedMap, dict]) -> Tuple[bool, str]:
         """
         测试站点是否可用
-        :param url: 站点域名
-        :param site_info: 站点信息
+        :param site_info: 站点域名或者站点的数据对象
         :return: (是否可用, 错误信息)
         """
         # 检查域名是否可用
-        if url:
+        if isinstance(site_info, str):
+            url = site_info
             domain = StringUtils.get_url_domain(url)
             site_info = self.siteoper.get_by_domain(domain)
-        else:
-            domain = StringUtils.get_url_domain(site_info.get("url"))
-
-        if not site_info:
-            return False, f"站点【{url}】不存在"
-
-        # 模拟登录
-        try:
-            # 特殊站点测试
-            if self.special_site_test.get(domain):
-                return self.special_site_test[domain](site_info)
-
+            if not site_info:
+                return False, f"站点【{url}】不存在"
             # 通用站点测试
             site_url = site_info.url
             site_cookie = site_info.cookie
@@ -358,7 +348,21 @@ class SiteChain(ChainBase):
             public = site_info.public
             proxies = settings.PROXY if site_info.proxy else None
             proxy_server = settings.PROXY_SERVER if site_info.proxy else None
-
+        else:
+            # 外部站点测试
+            site_url = site_info.get("url")
+            domain = StringUtils.get_url_domain(site_url)
+            site_cookie = site_info.get("cookie")
+            ua = site_info.get("ua")
+            render = site_info.get("render")
+            public = site_info.get("public")
+            proxies = settings.PROXY if site_info.get("proxy") else None
+            proxy_server = settings.PROXY_SERVER if site_info.get("proxy") else None
+        # 模拟登录
+        try:
+            # 特殊站点测试
+            if self.special_site_test.get(domain):
+                return self.special_site_test[domain](site_info)
             # 访问链接
             if render:
                 page_source = PlaywrightHelper().get_page_source(url=site_url,