pyrofork: Refactor test_mode

Signed-off-by: wulan17 <wulan17@komodos.id>

pyrogram/client.py

@@ -897,6 +897,8 @@ class Client(Methods):
                         except Exception as e:
                             print(e)
             # Needed for migration from storage v3 to v4
+            if self.in_memory or self.session_string:
+                await self.insert_default_dc_options()
             if not await self.storage.get_dc_address(await self.storage.dc_id(), self.ipv6):
                 log.info("No DC address found, inserting default DC options...")
                 await self.insert_default_dc_options()

pyrogram/storage/file_storage.py

@@ -38,19 +38,6 @@ CREATE TABLE update_state
 );
 """
 
-UPDATE_DC_SCHEMA = """
-CREATE TABLE dc_options
-(
-    id       INTEGER PRIMARY KEY AUTOINCREMENT,
-    dc_id    INTEGER,
-    address  TEXT,
-    port     INTEGER,
-    is_ipv6  BOOLEAN,
-    is_media BOOLEAN,
-    UNIQUE(dc_id, is_ipv6, is_media)
-);
-"""
-
 
 class FileStorage(SQLiteStorage):
     FILE_EXTENSION = ".session"
@@ -83,7 +70,7 @@ class FileStorage(SQLiteStorage):
 
         if version == 4:
             with self.conn:
-                self.conn.executescript(UPDATE_DC_SCHEMA)
+                self.conn.executescript(self.UPDATE_DC_SCHEMA)
 
             version += 1
 

pyrogram/storage/mongo_storage.py

@@ -241,11 +241,11 @@ class MongoStorage(Storage):
             return
 
         await self._dc_options.update_one(
-            {"$and": {
+            {"$and": [
                 {'dc_id': value[0]},
                 {'is_ipv6': value[3]},
-                {'is_media': value[4]},
-            }},
+                {'is_media': value[4]}
+            ]},
             {'$set': {'address': value[1], 'port': value[2]}},
             upsert=True
         )

pyrogram/storage/sqlite_storage.py

@@ -97,6 +97,20 @@ END;
 """
 
 
+UPDATE_DC_SCHEMA = """
+CREATE TABLE dc_options
+(
+    id       INTEGER PRIMARY KEY AUTOINCREMENT,
+    dc_id    INTEGER,
+    address  TEXT,
+    port     INTEGER,
+    is_ipv6  BOOLEAN,
+    is_media BOOLEAN,
+    UNIQUE(dc_id, is_ipv6, is_media)
+);
+"""
+
+
 def get_input_peer(peer_id: int, access_hash: int, peer_type: str):
     if peer_type in ["user", "bot"]:
         return raw.types.InputPeerUser(
@@ -121,6 +135,7 @@ def get_input_peer(peer_id: int, access_hash: int, peer_type: str):
 class SQLiteStorage(Storage):
     VERSION = 4
     USERNAME_TTL = 8 * 60 * 60
+    UPDATE_DC_SCHEMA = globals().get("UPDATE_DC_SCHEMA", "")
 
     def __init__(self, name: str):
         super().__init__(name)
@@ -131,6 +146,7 @@ class SQLiteStorage(Storage):
         with self.conn:
             self.conn.executescript(SCHEMA)
             self.conn.executescript(UNAME_SCHEMA)
+            self.conn.executescript(self.UPDATE_DC_SCHEMA)
 
             self.conn.execute(
                 "INSERT INTO version VALUES (?)",