Cleanup empty dir if frontend zip download failed (Comfy-Org#4574)

Author: huchenlei

app/frontend_management.py

@@ -8,7 +8,7 @@
 from dataclasses import dataclass
 from functools import cached_property
 from pathlib import Path
-from typing import TypedDict
+from typing import TypedDict, Optional
 
 import requests
 from typing_extensions import NotRequired
@@ -132,12 +132,13 @@ def parse_version_string(cls, value: str) -> tuple[str, str, str]:
         return match_result.group(1), match_result.group(2), match_result.group(3)
 
     @classmethod
-    def init_frontend_unsafe(cls, version_string: str) -> str:
+    def init_frontend_unsafe(cls, version_string: str, provider: Optional[FrontEndProvider] = None) -> str:
         """
         Initializes the frontend for the specified version.
 
         Args:
             version_string (str): The version string.
+            provider (FrontEndProvider, optional): The provider to use. Defaults to None.
 
         Returns:
             str: The path to the initialized frontend.
@@ -150,23 +151,29 @@ def init_frontend_unsafe(cls, version_string: str) -> str:
             return cls.DEFAULT_FRONTEND_PATH
 
         repo_owner, repo_name, version = cls.parse_version_string(version_string)
-        provider = FrontEndProvider(repo_owner, repo_name)
+        provider = provider or FrontEndProvider(repo_owner, repo_name)
         release = provider.get_release(version)
 
         semantic_version = release["tag_name"].lstrip("v")
         web_root = str(
             Path(cls.CUSTOM_FRONTENDS_ROOT) / provider.folder_name / semantic_version
         )
         if not os.path.exists(web_root):
-            os.makedirs(web_root, exist_ok=True)
-            logging.info(
-                "Downloading frontend(%s) version(%s) to (%s)",
-                provider.folder_name,
-                semantic_version,
-                web_root,
-            )
-            logging.debug(release)
-            download_release_asset_zip(release, destination_path=web_root)
+            try:
+                os.makedirs(web_root, exist_ok=True)
+                logging.info(
+                    "Downloading frontend(%s) version(%s) to (%s)",
+                    provider.folder_name,
+                    semantic_version,
+                    web_root,
+                )
+                logging.debug(release)
+                download_release_asset_zip(release, destination_path=web_root)
+            finally:
+                # Clean up the directory if it is empty, i.e. the download failed
+                if not os.listdir(web_root):
+                    os.rmdir(web_root)
+
         return web_root
 
     @classmethod

tests-unit/app_test/frontend_manager_test.py

@@ -1,6 +1,7 @@
 import argparse
 import pytest
 from requests.exceptions import HTTPError
+from unittest.mock import patch
 
 from app.frontend_management import (
     FrontendManager,
@@ -83,6 +84,35 @@ def test_init_frontend_invalid_provider():
     with pytest.raises(HTTPError):
         FrontendManager.init_frontend_unsafe(version_string)
 
+@pytest.fixture
+def mock_os_functions():
+    with patch('app.frontend_management.os.makedirs') as mock_makedirs, \
+         patch('app.frontend_management.os.listdir') as mock_listdir, \
+         patch('app.frontend_management.os.rmdir') as mock_rmdir:
+        mock_listdir.return_value = []  # Simulate empty directory
+        yield mock_makedirs, mock_listdir, mock_rmdir
+
+@pytest.fixture
+def mock_download():
+    with patch('app.frontend_management.download_release_asset_zip') as mock:
+        mock.side_effect = Exception("Download failed")  # Simulate download failure
+        yield mock
+
+def test_finally_block(mock_os_functions, mock_download, mock_provider):
+    # Arrange
+    mock_makedirs, mock_listdir, mock_rmdir = mock_os_functions
+    version_string = 'test-owner/[email protected]'
+
+    # Act & Assert
+    with pytest.raises(Exception):
+        FrontendManager.init_frontend_unsafe(version_string, mock_provider)
+
+    # Assert
+    mock_makedirs.assert_called_once()
+    mock_download.assert_called_once()
+    mock_listdir.assert_called_once()
+    mock_rmdir.assert_called_once()
+
 
 def test_parse_version_string():
     version_string = "owner/[email protected]"