59 lines
2.0 KiB
Python
59 lines
2.0 KiB
Python
import sys
|
|
import types
|
|
import unittest
|
|
from unittest.mock import Mock, patch
|
|
|
|
from comfyui_to_python.node_runtime import cleanup_comfyui_runtime
|
|
|
|
|
|
class NodeRuntimeCleanupTest(unittest.TestCase):
|
|
def test_cleanup_releases_caches_without_forcing_model_unload_by_default(self):
|
|
comfy_module = types.ModuleType("comfy")
|
|
comfy_module.__path__ = []
|
|
model_management = types.ModuleType("comfy.model_management")
|
|
model_management.cleanup_models_gc = Mock()
|
|
model_management.unload_all_models = Mock()
|
|
model_management.soft_empty_cache = Mock()
|
|
comfy_module.model_management = model_management
|
|
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"comfy": comfy_module,
|
|
"comfy.model_management": model_management,
|
|
},
|
|
), patch.dict("os.environ", {}, clear=False):
|
|
cleanup_comfyui_runtime()
|
|
|
|
model_management.cleanup_models_gc.assert_called_once_with()
|
|
model_management.soft_empty_cache.assert_called_once_with()
|
|
model_management.unload_all_models.assert_not_called()
|
|
|
|
def test_cleanup_can_force_model_unload_from_environment(self):
|
|
comfy_module = types.ModuleType("comfy")
|
|
comfy_module.__path__ = []
|
|
model_management = types.ModuleType("comfy.model_management")
|
|
model_management.cleanup_models_gc = Mock()
|
|
model_management.unload_all_models = Mock()
|
|
model_management.soft_empty_cache = Mock()
|
|
comfy_module.model_management = model_management
|
|
|
|
with patch.dict(
|
|
sys.modules,
|
|
{
|
|
"comfy": comfy_module,
|
|
"comfy.model_management": model_management,
|
|
},
|
|
), patch.dict(
|
|
"os.environ",
|
|
{"COMFYUI_TOPYTHON_UNLOAD_MODELS": "true"},
|
|
clear=False,
|
|
):
|
|
cleanup_comfyui_runtime()
|
|
|
|
model_management.unload_all_models.assert_called_once_with()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|