🎨 catch more exception.

addons/builtin/tiktok.py

@@ -1,48 +1,41 @@
 import asyncio
 
 import aiohttp
-from aiohttp import ClientConnectorError
+import async_timeout
 from loguru import logger
-
 from utils import retry
 
 
 # collector section
-@retry(2)
+@retry(1)
 async def fetch_tiktok(collector, session: aiohttp.ClientSession, proxy=None):
     """
-    tiktok解锁测试,检查测速是否被ban
+    tiktok解锁测试
     :param collector:
     :param session:
     :param proxy:
     :return:
     """
     tiktokurl = 'https://www.tiktok.com'
-    try:
-        # async with aiohttp.ClientSession() as session:
-        async with session.get(tiktokurl, proxy=proxy, timeout=5) as resq:
-            if resq.status == 200:
-                response_text = await resq.text()
-                region = response_text.find('"region":')
-                if region != -1:
-                    region = response_text[region:].split('"')[3]
-                    # print("Tiktok Region: ", region)
-                    collector.info['tiktok'] = f"解锁({region})"
+    conn = session.connector
+    if type(conn).__name__ == 'ProxyConnector':
+        proxy = "http://" + conn._proxy_host + ":" + str(conn._proxy_port)
+    async with async_timeout.timeout(10):
+        async with aiohttp.ClientSession() as session:
+            async with session.get(tiktokurl, proxy=proxy, timeout=5) as resp:
+                if resp.status == 200:
+                    response_text = await resp.text()
+                    region = response_text.find('"region":')
+                    if region != -1:
+                        region = response_text[region:].split('"')[3]
+                        # print("Tiktok Region: ", region)
+                        collector.info['tiktok'] = f"解锁({region})"
+                    else:
+                        # print("Tiktok Region: Not found")
+                        collector.info['tiktok'] = "失败"
                 else:
-                    # print("Tiktok Region: Not found")
-                    collector.info['tiktok'] = "失败"
-            else:
-                collector.info['tiktok'] = "未知"
-        return True
-    except ClientConnectorError as c:
-        logger.warning("tiktok请求发生错误:" + str(c))
-        await fetch_tiktok(collector, session=session, proxy=proxy)
-        collector.info['tiktok'] = "连接错误"
-        return False
-    except asyncio.exceptions.TimeoutError:
-        await fetch_tiktok(collector, session=session, proxy=proxy)
-        collector.info['tiktok'] = "超时"
-        return False
+                    collector.info['tiktok'] = "未知"
+    return True
 
 
 def task(Collector, session, proxy):
@@ -57,8 +50,10 @@ def get_tiktok_info(ReCleaner):
     """
     try:
         if 'tiktok' not in ReCleaner.data:
+            logger.warning("采集器内无数据")
             return "N/A"
         else:
+            # logger.info("tiktok解锁：" + str(ReCleaner.data.get('tiktok', "N/A")))
             return ReCleaner.data.get('tiktok', "N/A")
     except Exception as e:
         logger.error(e)
@@ -71,13 +66,23 @@ def get_tiktok_info(ReCleaner):
     "GET": get_tiktok_info
 }
 
+if __name__ == "__main__":
+    async def demo():
+        class FakeColl:
+            def __init__(self):
+                self.info = {}
+                self.data = self.info
 
-async def demo():
-    from utils import script_demo
-    await script_demo(fetch_tiktok, proxy='http://127.0.0.1:11112')
+        fakecl = FakeColl()
+        from aiohttp_socks import ProxyConnector
+        conn = ProxyConnector(host="127.0.0.1", port=11112, limit=0)
+        session = aiohttp.ClientSession(connector=conn)
+        await fetch_tiktok(fakecl, session, proxy=None)
+        print(fakecl.info)
+        await session.close()
 
 
-if __name__ == "__main__":
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    loop.run_until_complete(demo())
+    if __name__ == "__main__":
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        loop.run_until_complete(demo())