mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-01-10 14:20:49 +08:00
60 lines
2.0 KiB
Python
60 lines
2.0 KiB
Python
import os
|
|
import tempfile
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
|
|
from comfy_extras.nodes.nodes_language import SaveString
|
|
|
|
|
|
@pytest.fixture
|
|
def save_string_node():
|
|
return SaveString()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_get_save_path(save_string_node):
|
|
with patch.object(save_string_node, 'get_save_path') as mock_method:
|
|
mock_method.return_value = (tempfile.gettempdir(), "test", 0, "", "test")
|
|
yield mock_method
|
|
|
|
|
|
def test_save_string_single(save_string_node, mock_get_save_path):
|
|
test_string = "Test string content"
|
|
result = save_string_node.execute(test_string, "test_prefix", ".txt")
|
|
|
|
assert result == {"ui": {"string": [test_string]}}
|
|
mock_get_save_path.assert_called_once_with("test_prefix")
|
|
|
|
saved_file_path = os.path.join(tempfile.gettempdir(), "test_00000_.txt")
|
|
assert os.path.exists(saved_file_path)
|
|
with open(saved_file_path, "r") as f:
|
|
assert f.read() == test_string
|
|
|
|
|
|
def test_save_string_list(save_string_node, mock_get_save_path):
|
|
test_strings = ["First string", "Second string", "Third string"]
|
|
result = save_string_node.execute(test_strings, "test_prefix", ".txt")
|
|
|
|
assert result == {"ui": {"string": test_strings}}
|
|
mock_get_save_path.assert_called_once_with("test_prefix")
|
|
|
|
for i, test_string in enumerate(test_strings):
|
|
saved_file_path = os.path.join(tempfile.gettempdir(), f"test_00000_{i:02d}_.txt")
|
|
assert os.path.exists(saved_file_path)
|
|
with open(saved_file_path, "r") as f:
|
|
assert f.read() == test_string
|
|
|
|
|
|
def test_save_string_default_extension(save_string_node, mock_get_save_path):
|
|
test_string = "Test string content"
|
|
result = save_string_node.execute(test_string, "test_prefix")
|
|
|
|
assert result == {"ui": {"string": [test_string]}}
|
|
mock_get_save_path.assert_called_once_with("test_prefix")
|
|
|
|
saved_file_path = os.path.join(tempfile.gettempdir(), "test_00000_.json")
|
|
assert os.path.exists(saved_file_path)
|
|
with open(saved_file_path, "r") as f:
|
|
assert f.read() == test_string
|