From 37c1052b390e63750b37c753b7e357d214366a75 Mon Sep 17 00:00:00 2001
From: Vftdan <vftdan@gmail.com>
Date: Tue, 11 Jun 2024 20:01:24 +0200
Subject: [PATCH] Add connect-timeout config option

---
 config.ini              | 4 ++++
 mastoposter/__main__.py | 1 +
 mastoposter/sources.py  | 5 +++--
 3 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/config.ini b/config.ini
index f25820e..0ebd7a5 100644
--- a/config.ini
+++ b/config.ini
@@ -42,6 +42,10 @@ list = 1
 auto-reconnect = yes
 reconnect-delay = 1.0
 
+# Change websocket connection opening timeout.
+# It may be useful when initial server connection may take a long time.
+connect-timeout = 60.0
+
 # Number of retries in case request fails. Applies globally
 # Can be changed on per-module basis
 http-retries = 5
diff --git a/mastoposter/__main__.py b/mastoposter/__main__.py
index 5bc19ac..682e3c1 100644
--- a/mastoposter/__main__.py
+++ b/mastoposter/__main__.py
@@ -155,6 +155,7 @@ def main():
             ),
             reconnect=conf["main"].getboolean("auto_reconnect", False),
             reconnect_delay=conf["main"].getfloat("reconnect_delay", 1.0),
+            connect_timeout=conf["main"].getfloat("connect_timeout", 60.0),
             list=conf["main"]["list"],
             access_token=conf["main"]["token"],
         )
diff --git a/mastoposter/sources.py b/mastoposter/sources.py
index 9ef4478..39ad42b 100644
--- a/mastoposter/sources.py
+++ b/mastoposter/sources.py
@@ -24,7 +24,8 @@ logger = getLogger("sources")
 
 
 async def websocket_source(
-    url: str, reconnect: bool = False, reconnect_delay: float = 1.0, **params
+    url: str, reconnect: bool = False, reconnect_delay: float = 1.0,
+    connect_timeout: float = 60.0, **params
 ) -> AsyncGenerator[Status, None]:
     from websockets.client import connect
     from websockets.exceptions import WebSocketException
@@ -37,7 +38,7 @@ async def websocket_source(
     while True:
         try:
             logger.info("attempting to connect to %s", public_url)
-            async with connect(url, open_timeout=60) as ws:
+            async with connect(url, open_timeout=connect_timeout) as ws:
                 logger.info("Connected to WebSocket")
                 while (msg := await ws.recv()) is not None:
                     event = loads(msg)