fix: add timeout to PyPI update check to prevent hanging offline

comfy_cli/update.py

@@ -1,3 +1,4 @@
+import logging
 import sys
 from importlib.metadata import metadata
 
@@ -6,22 +7,23 @@
 from rich.console import Console
 from rich.panel import Panel
 
+logger = logging.getLogger(__name__)
 console = Console()
 
 
 def check_for_newer_pypi_version(package_name, current_version):
     url = f"https://pypi.org/pypi/{package_name}/json"
     try:
-        response = requests.get(url)
-        response.raise_for_status()  # Raises stored HTTPError, if one occurred
+        response = requests.get(url, timeout=5)
+        response.raise_for_status()
         latest_version = response.json()["info"]["version"]
 
         if version.parse(latest_version) > version.parse(current_version):
             return True, latest_version
 
         return False, current_version
-    except requests.RequestException:
-        # print(f"Error checking latest version: {e}")
+    except requests.RequestException as e:
+        logger.warning(f"Failed to check for updates: {e}")
         return False, current_version
 
 

tests/comfy_cli/test_update.py

@@ -0,0 +1,59 @@
+from unittest.mock import MagicMock, patch
+
+import requests
+
+from comfy_cli.update import check_for_newer_pypi_version, check_for_updates
+
+
+def _mock_pypi_response(latest_version):
+    mock_resp = MagicMock()
+    mock_resp.status_code = 200
+    mock_resp.json.return_value = {"info": {"version": latest_version}}
+    return mock_resp
+
+
+class TestCheckForNewerPypiVersion:
+    @patch("comfy_cli.update.requests.get")
+    def test_newer_version_available(self, mock_get):
+        mock_get.return_value = _mock_pypi_response("99.0.0")
+        has_newer, ver = check_for_newer_pypi_version("comfy-cli", "1.0.0")
+        assert has_newer is True
+        assert ver == "99.0.0"
+
+    @patch("comfy_cli.update.requests.get")
+    def test_no_update_when_current(self, mock_get):
+        mock_get.return_value = _mock_pypi_response("1.0.0")
+        has_newer, ver = check_for_newer_pypi_version("comfy-cli", "1.0.0")
+        assert has_newer is False
+        assert ver == "1.0.0"
+
+    @patch("comfy_cli.update.requests.get")
+    def test_network_failure_returns_false(self, mock_get):
+        mock_get.side_effect = requests.Timeout("connection timed out")
+        has_newer, ver = check_for_newer_pypi_version("comfy-cli", "1.0.0")
+        assert has_newer is False
+        assert ver == "1.0.0"
+
+    @patch("comfy_cli.update.requests.get")
+    def test_timeout_value_is_passed(self, mock_get):
+        mock_get.return_value = _mock_pypi_response("1.0.0")
+        check_for_newer_pypi_version("comfy-cli", "1.0.0")
+        mock_get.assert_called_once_with("https://pypi.org/pypi/comfy-cli/json", timeout=5)
+
+
+class TestCheckForUpdates:
+    @patch("comfy_cli.update.notify_update")
+    @patch("comfy_cli.update.get_version_from_pyproject", return_value="1.0.0")
+    @patch("comfy_cli.update.requests.get")
+    def test_notifies_when_update_available(self, mock_get, _mock_ver, mock_notify):
+        mock_get.return_value = _mock_pypi_response("2.0.0")
+        check_for_updates()
+        mock_notify.assert_called_once_with("1.0.0", "2.0.0")
+
+    @patch("comfy_cli.update.notify_update")
+    @patch("comfy_cli.update.get_version_from_pyproject", return_value="1.0.0")
+    @patch("comfy_cli.update.requests.get")
+    def test_no_notification_on_network_error(self, mock_get, _mock_ver, mock_notify):
+        mock_get.side_effect = requests.ConnectionError("offline")
+        check_for_updates()
+        mock_notify.assert_not_called()