cbpark84
1c919f1106
- run_batch(): sequential|concurrent|auto - chain(): 파이프라인 체이닝 - schema(): 스키마 자동 추출 - ResultCache: LRU+TTL 캐시 - preview_callback: 실시간 프리뷰 - upload_file/audio/video() - BridgeInputSEED: random/increment/fixed - MCPServer: MCP Tool 자동 생성
318 lines
12 KiB
Python
318 lines
12 KiB
Python
"""
|
||
nodes/server_routes.py 유닛 테스트.
|
||
ComfyUI 없이 실행 가능 — register_routes()는 호출하지 않음.
|
||
"""
|
||
|
||
import base64
|
||
import io
|
||
import json
|
||
import struct
|
||
import pytest
|
||
import numpy as np
|
||
from PIL import Image
|
||
from unittest.mock import patch, MagicMock
|
||
|
||
|
||
# ── 임포트 ────────────────────────────────────────────────────
|
||
from nodes.server_routes import (
|
||
_numpy_to_wav,
|
||
_encode_pil_image,
|
||
_encode_audio,
|
||
_encode_latent,
|
||
_encode_conditioning,
|
||
_encode_comfy_object,
|
||
_encode_value,
|
||
store_bridge_result,
|
||
_results_store,
|
||
_binary_store,
|
||
_INLINE_THRESHOLD,
|
||
)
|
||
|
||
|
||
# ── WAV 인코딩 테스트 ─────────────────────────────────────────
|
||
|
||
class TestNumpyToWav:
|
||
"""_numpy_to_wav 함수: numpy [C, T] → WAV bytes"""
|
||
|
||
def test_mono_produces_valid_riff_header(self):
|
||
audio = np.zeros((1, 44100), dtype=np.float32) # 1초 모노
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
|
||
assert wav[:4] == b"RIFF"
|
||
assert wav[8:12] == b"WAVE"
|
||
assert wav[12:16] == b"fmt "
|
||
assert wav[36:40] == b"data"
|
||
|
||
def test_stereo_channel_count_in_header(self):
|
||
audio = np.zeros((2, 22050), dtype=np.float32) # 스테레오
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
|
||
# NumChannels at offset 22 (little-endian uint16)
|
||
num_channels = struct.unpack_from("<H", wav, 22)[0]
|
||
assert num_channels == 2
|
||
|
||
def test_sample_rate_in_header(self):
|
||
audio = np.zeros((1, 16000), dtype=np.float32)
|
||
wav = _numpy_to_wav(audio, sample_rate=16000)
|
||
|
||
sample_rate = struct.unpack_from("<I", wav, 24)[0]
|
||
assert sample_rate == 16000
|
||
|
||
def test_data_size_matches(self):
|
||
samples = 1000
|
||
audio = np.zeros((1, samples), dtype=np.float32)
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
|
||
# data chunk size at offset 40
|
||
data_size = struct.unpack_from("<I", wav, 40)[0]
|
||
assert data_size == samples * 1 * 2 # samples × channels × 2 bytes
|
||
|
||
def test_silence_produces_zero_data(self):
|
||
audio = np.zeros((1, 100), dtype=np.float32)
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
data = wav[44:] # after 44-byte header
|
||
assert all(b == 0 for b in data)
|
||
|
||
def test_clipping_positive(self):
|
||
"""1.0 이상 값 → int16 max (32767)"""
|
||
audio = np.ones((1, 1), dtype=np.float32) * 2.0
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
sample = struct.unpack_from("<h", wav, 44)[0]
|
||
assert sample == 32767
|
||
|
||
def test_clipping_negative(self):
|
||
"""-1.0 이하 값 → int16 min (-32768)"""
|
||
audio = np.ones((1, 1), dtype=np.float32) * -2.0
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
sample = struct.unpack_from("<h", wav, 44)[0]
|
||
assert sample == -32768
|
||
|
||
def test_1d_input_treated_as_mono(self):
|
||
"""1D 입력도 처리 가능해야 함"""
|
||
audio = np.zeros(100, dtype=np.float32)
|
||
wav = _numpy_to_wav(audio, sample_rate=44100)
|
||
assert wav[:4] == b"RIFF"
|
||
|
||
|
||
# ── 이미지 인코딩 테스트 ──────────────────────────────────────
|
||
|
||
class TestEncodePilImage:
|
||
"""_encode_pil_image 함수: PIL Image → meta + binary"""
|
||
|
||
def _make_image(self, width=64, height=64, color=(128, 64, 32)):
|
||
img = Image.new("RGB", (width, height), color)
|
||
return img
|
||
|
||
def test_small_image_has_inline_data(self):
|
||
img = self._make_image(64, 64)
|
||
meta, binary = _encode_pil_image(img, "/comfybridge/file/pid/name")
|
||
assert "data" in meta
|
||
assert meta["__type__"] == "image"
|
||
assert meta["width"] == 64
|
||
assert meta["height"] == 64
|
||
|
||
def test_small_image_data_is_valid_png(self):
|
||
img = self._make_image(64, 64)
|
||
meta, binary = _encode_pil_image(img, "/comfybridge/file/pid/name")
|
||
png_bytes = base64.b64decode(meta["data"])
|
||
recovered = Image.open(io.BytesIO(png_bytes))
|
||
assert recovered.size == (64, 64)
|
||
|
||
def test_small_image_binary_is_png(self):
|
||
img = self._make_image(64, 64)
|
||
meta, (raw_bytes, content_type) = _encode_pil_image(img, "/url")
|
||
assert content_type == "image/png"
|
||
assert raw_bytes[:8] == b"\x89PNG\r\n\x1a\n" # PNG magic
|
||
|
||
def test_small_image_url_always_present(self):
|
||
img = self._make_image(64, 64)
|
||
meta, _ = _encode_pil_image(img, "/comfybridge/file/pid/name")
|
||
assert meta["url"] == "/comfybridge/file/pid/name"
|
||
|
||
def test_large_image_no_inline_data(self):
|
||
"""2MB 이상 이미지는 inline base64 생략"""
|
||
img = self._make_image(64, 64)
|
||
# _INLINE_THRESHOLD보다 큰 PNG를 흉내내기 위해 임계값을 낮춤
|
||
with patch("nodes.server_routes._INLINE_THRESHOLD", 1):
|
||
meta, binary = _encode_pil_image(img, "/url")
|
||
assert "data" not in meta
|
||
assert "url" in meta
|
||
|
||
def test_large_image_has_url(self):
|
||
img = self._make_image(64, 64)
|
||
with patch("nodes.server_routes._INLINE_THRESHOLD", 1):
|
||
meta, _ = _encode_pil_image(img, "/comfybridge/file/pid/large")
|
||
assert meta["url"] == "/comfybridge/file/pid/large"
|
||
|
||
def test_size_bytes_reported(self):
|
||
img = self._make_image(64, 64)
|
||
meta, (raw_bytes, _) = _encode_pil_image(img, "/url")
|
||
assert meta["size_bytes"] == len(raw_bytes)
|
||
|
||
|
||
# ── 오디오 인코딩 테스트 ─────────────────────────────────────
|
||
|
||
class TestEncodeAudio:
|
||
"""_encode_audio 함수: {"waveform": ..., "sample_rate": int} → meta + WAV"""
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def require_torch(self):
|
||
pytest.importorskip("torch")
|
||
|
||
def _make_audio_dict(self, channels=1, samples=44100, sample_rate=44100):
|
||
try:
|
||
import torch
|
||
waveform = torch.zeros(1, channels, samples)
|
||
except ImportError:
|
||
waveform = np.zeros((1, channels, samples), dtype=np.float32)
|
||
return {"waveform": waveform, "sample_rate": sample_rate}
|
||
|
||
def test_meta_type_is_audio(self):
|
||
value = self._make_audio_dict()
|
||
meta, binary = _encode_audio(value, "/url")
|
||
assert meta["__type__"] == "audio"
|
||
|
||
def test_meta_has_url(self):
|
||
value = self._make_audio_dict()
|
||
meta, _ = _encode_audio(value, "/comfybridge/file/p/a")
|
||
assert meta["url"] == "/comfybridge/file/p/a"
|
||
|
||
def test_meta_format_is_wav(self):
|
||
value = self._make_audio_dict()
|
||
meta, _ = _encode_audio(value, "/url")
|
||
assert meta["format"] == "WAV"
|
||
|
||
def test_meta_duration(self):
|
||
value = self._make_audio_dict(samples=44100, sample_rate=44100)
|
||
meta, _ = _encode_audio(value, "/url")
|
||
assert abs(meta["duration"] - 1.0) < 0.01
|
||
|
||
def test_binary_content_type(self):
|
||
value = self._make_audio_dict()
|
||
_, (raw, ct) = _encode_audio(value, "/url")
|
||
assert ct == "audio/wav"
|
||
|
||
def test_binary_is_valid_wav(self):
|
||
value = self._make_audio_dict()
|
||
_, (raw, _) = _encode_audio(value, "/url")
|
||
assert raw[:4] == b"RIFF"
|
||
assert raw[8:12] == b"WAVE"
|
||
|
||
|
||
# ── Latent 인코딩 테스트 ─────────────────────────────────────
|
||
|
||
class TestEncodeLatent:
|
||
def test_latent_meta_with_torch(self):
|
||
pytest.importorskip("torch")
|
||
import torch
|
||
value = {"samples": torch.zeros(1, 4, 64, 64)}
|
||
meta = _encode_latent(value)
|
||
assert meta["__type__"] == "latent"
|
||
assert meta["batch_size"] == 1
|
||
assert meta["channels"] == 4
|
||
assert meta["height_latent"] == 64
|
||
assert meta["width_latent"] == 64
|
||
assert meta["estimated_pixel_height"] == 512
|
||
assert meta["estimated_pixel_width"] == 512
|
||
|
||
def test_latent_fallback_without_samples(self):
|
||
value = {"other_key": "something"}
|
||
meta = _encode_latent(value)
|
||
assert meta["__type__"] == "latent"
|
||
assert "keys" in meta
|
||
|
||
|
||
# ── Conditioning 인코딩 테스트 ───────────────────────────────
|
||
|
||
class TestEncodeConditioning:
|
||
def test_conditioning_structure(self):
|
||
pytest.importorskip("torch")
|
||
import torch
|
||
cond = [[torch.zeros(1, 77, 768), {"pooled_output": torch.zeros(1, 1280)}]]
|
||
meta = _encode_conditioning(cond)
|
||
assert meta["__type__"] == "conditioning"
|
||
assert len(meta["conditions"]) == 1
|
||
assert meta["conditions"][0]["shape"] == [1, 77, 768]
|
||
|
||
def test_conditioning_empty_list(self):
|
||
meta = _encode_conditioning([])
|
||
assert meta["__type__"] == "conditioning"
|
||
|
||
|
||
# ── _encode_value 통합 테스트 ─────────────────────────────────
|
||
|
||
class TestEncodeValue:
|
||
def test_scalar_string(self):
|
||
meta, binary = _encode_value("p1", "n1", "hello")
|
||
assert meta == "hello"
|
||
assert binary is None
|
||
|
||
def test_scalar_int(self):
|
||
meta, binary = _encode_value("p1", "n1", 42)
|
||
assert meta == 42
|
||
assert binary is None
|
||
|
||
def test_scalar_float(self):
|
||
meta, binary = _encode_value("p1", "n1", 3.14)
|
||
assert abs(meta - 3.14) < 1e-6
|
||
assert binary is None
|
||
|
||
def test_scalar_none(self):
|
||
meta, binary = _encode_value("p1", "n1", None)
|
||
assert meta is None
|
||
assert binary is None
|
||
|
||
def test_pil_image(self):
|
||
img = Image.new("RGB", (32, 32), (255, 0, 0))
|
||
meta, binary = _encode_value("p1", "n1", img)
|
||
assert meta["__type__"] == "image"
|
||
assert binary is not None
|
||
assert binary[1] == "image/png"
|
||
|
||
def test_audio_dict(self):
|
||
try:
|
||
import torch
|
||
waveform = torch.zeros(1, 1, 1000)
|
||
except ImportError:
|
||
pytest.skip("torch not available")
|
||
value = {"waveform": waveform, "sample_rate": 44100}
|
||
meta, binary = _encode_value("p1", "n1", value)
|
||
assert meta["__type__"] == "audio"
|
||
assert binary is not None
|
||
|
||
|
||
# ── store_bridge_result 테스트 ───────────────────────────────
|
||
|
||
class TestStoreBridgeResult:
|
||
def setup_method(self):
|
||
"""각 테스트 전 저장소 초기화"""
|
||
_results_store.clear()
|
||
_binary_store.clear()
|
||
|
||
def test_scalar_stored_in_results(self):
|
||
store_bridge_result("pid1", "caption", "hello world")
|
||
assert _results_store["pid1"]["caption"] == "hello world"
|
||
# 스칼라는 binary_store에 없어야 함
|
||
assert "caption" not in _binary_store.get("pid1", {})
|
||
|
||
def test_image_stored_in_both(self):
|
||
img = Image.new("RGB", (32, 32), (0, 255, 0))
|
||
store_bridge_result("pid1", "out_img", img)
|
||
assert _results_store["pid1"]["out_img"]["__type__"] == "image"
|
||
assert "out_img" in _binary_store["pid1"]
|
||
|
||
def test_multiple_outputs_same_prompt(self):
|
||
store_bridge_result("pid1", "caption", "text")
|
||
store_bridge_result("pid1", "score", 0.95)
|
||
assert "caption" in _results_store["pid1"]
|
||
assert "score" in _results_store["pid1"]
|
||
|
||
def test_lru_eviction(self):
|
||
"""MAX_ENTRIES 초과 시 가장 오래된 항목 제거"""
|
||
with patch("nodes.server_routes._MAX_ENTRIES", 3):
|
||
for i in range(4):
|
||
store_bridge_result(f"pid{i}", "val", i)
|
||
# pid0이 제거됐는지 확인 (LRU)
|
||
assert "pid0" not in _results_store
|
||
assert "pid3" in _results_store
|