Coverage for tests/test_wrapper_wan22.py: 100%
175 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import sys
4import gc
5import pytest
7from unittest.mock import patch
8from unittest.mock import MagicMock
9from unittest.mock import AsyncMock
11from PIL import Image
13from tests.torch_mock import TorchMock
15mock_torch = TorchMock()
17sys.path.append("wrapper")
18sys.path.append("wrapper/wan")
19sys.path.append("wrapper/wan22")
21# Build a mock for media_utils with async base64_to_audio_file
22mock_media_utils = MagicMock()
23mock_media_utils.base64_to_audio_file = AsyncMock(return_value="/tmp/test_audio.wav")
24mock_media_utils.empty_audio_file = MagicMock(return_value="/tmp/warmup.wav")
26mock_modules = {
27 'nvidia_smi': MagicMock(),
28 'imageio': MagicMock(),
29 'cv2': MagicMock(),
30 'torch': mock_torch,
31 'torchvision': MagicMock(),
32 'torchvision.transforms': MagicMock(),
33 'torchvision.transforms.functional': MagicMock(),
34 'xfuser': MagicMock(),
35 'xfuser.config': MagicMock(),
36 'xfuser.core': MagicMock(),
37 'xfuser.core.distributed': MagicMock(),
38 'transformers': MagicMock(),
39 'wan': MagicMock(),
40 'wan.configs': MagicMock(),
41 'wan.modules': MagicMock(),
42 'wan.modules.t5': MagicMock(),
43 'wan.modules.clip': MagicMock(),
44 'wan.modules.vae': MagicMock(),
45 'wan.modules.model': MagicMock(),
46 'wan.utils': MagicMock(),
47 'wan.utils.utils': MagicMock(),
48 'wan.utils.fm_solvers_unipc': MagicMock(),
49 'wan.distributed': MagicMock(),
50 'wan.distributed.fsdp': MagicMock(),
51 'wan.distributed.sequence_parallel': MagicMock(),
52 'media_utils': mock_media_utils,
53}
54mock_modules.update(mock_torch.get_sub_modules())
56with patch.dict(sys.modules, mock_modules):
57 from image_utils import img_to_base64
58 from wan22.wrapper_wan22 import Wan22VideoGeneration
61@pytest.mark.asyncio
62async def test_init() -> None:
63 """Test basic initialization and health checks."""
64 model = Wan22VideoGeneration()
65 assert model is not None
66 assert model.model_name == "wan22"
68 model.init()
69 health = model.get_health()
70 assert health is not None
71 timestamps = model.get_timestamps()
72 assert timestamps is not None
74 del model
75 gc.collect()
78@pytest.mark.asyncio
79async def test_get_rest_args_missing_body() -> None:
80 """Test that get_rest_args raises for missing JSON body."""
81 model = Wan22VideoGeneration()
82 model.init()
84 with pytest.raises(ValueError, match="Missing JSON body"):
85 await model.get_rest_args(None)
88@pytest.mark.asyncio
89async def test_get_rest_args_missing_img() -> None:
90 """Test that get_rest_args raises when img is absent."""
91 model = Wan22VideoGeneration()
92 model.init()
94 with pytest.raises(ValueError, match="Missing 'img' parameter"):
95 await model.get_rest_args({})
98@pytest.mark.asyncio
99async def test_get_rest_args_missing_prompt() -> None:
100 """Test that get_rest_args raises when prompt is absent."""
101 model = Wan22VideoGeneration()
102 model.init()
104 img = Image.new("RGB", (40, 30))
105 img_base64 = img_to_base64(img)
107 with pytest.raises(ValueError, match="Missing 'prompt' parameter"):
108 await model.get_rest_args({"img": img_base64})
111@pytest.mark.asyncio
112async def test_get_rest_args_missing_audio() -> None:
113 """Test that get_rest_args raises when audio is absent and TTS is disabled."""
114 model = Wan22VideoGeneration()
115 model.init()
117 img = Image.new("RGB", (40, 30))
118 img_base64 = img_to_base64(img)
120 with pytest.raises(ValueError, match="Missing 'audio' parameter"):
121 await model.get_rest_args({
122 "img": img_base64,
123 "prompt": "test prompt",
124 })
127@pytest.mark.asyncio
128async def test_get_rest_args_with_audio() -> None:
129 """Test get_rest_args with a base64-encoded audio file."""
130 model = Wan22VideoGeneration()
131 model.init()
133 img = Image.new("RGB", (40, 30))
134 img_base64 = img_to_base64(img)
136 result = await model.get_rest_args({
137 "img": img_base64,
138 "prompt": "test prompt",
139 "audio": "dGVzdA==", # base64("test") - placeholder
140 })
142 assert "args" in result
143 assert result["task"] == "wan22"
144 args = result["args"]
145 assert args["prompt"] == "test prompt"
146 assert args["neg_prompt"] == ""
147 assert args["enable_tts"] is False
148 assert args["audio_path"] is not None
149 assert args["audio_path"].endswith(".wav")
150 assert args["max_area"] == 1024 * 704
151 assert args["sampling_steps"] == 40
152 assert args["infer_frames"] == 80
153 assert args["num_clip"] is None
155 del model
156 gc.collect()
159@pytest.mark.asyncio
160async def test_get_rest_args_with_tts() -> None:
161 """Test get_rest_args with TTS parameters."""
162 model = Wan22VideoGeneration()
163 model.init()
165 img = Image.new("RGB", (40, 30))
166 img_base64 = img_to_base64(img)
168 result = await model.get_rest_args({
169 "img": img_base64,
170 "prompt": "Summer beach scene",
171 "enable_tts": True,
172 "tts_prompt_audio": "dGVzdA==",
173 "tts_prompt_text": "Hello world",
174 "tts_text": "Generated speech text",
175 })
177 args = result["args"]
178 assert args["enable_tts"] is True
179 assert args["tts_text"] == "Generated speech text"
180 assert args["tts_prompt_text"] == "Hello world"
181 assert args["tts_prompt_audio"] is not None
182 assert args["tts_prompt_audio"].endswith(".wav")
183 assert args["audio_path"] is None
185 del model
186 gc.collect()
189@pytest.mark.asyncio
190async def test_get_rest_args_custom_params() -> None:
191 """Test get_rest_args honours custom resolution, steps, and clip count."""
192 model = Wan22VideoGeneration()
193 model.init()
195 img = Image.new("RGB", (40, 30))
196 img_base64 = img_to_base64(img)
198 result = await model.get_rest_args({
199 "img": img_base64,
200 "prompt": "test",
201 "audio": "dGVzdA==",
202 "max_area": 720 * 1280,
203 "sampling_steps": 20,
204 "infer_frames": 48,
205 "num_clip": 3,
206 "neg_prompt": "blurry",
207 "output_type": "video_path",
208 })
210 args = result["args"]
211 assert args["max_area"] == 720 * 1280
212 assert args["sampling_steps"] == 20
213 assert args["infer_frames"] == 48
214 assert args["num_clip"] == 3
215 assert args["neg_prompt"] == "blurry"
216 assert args["output_type"] == "video_path"
218 del model
219 gc.collect()
222@pytest.mark.asyncio
223async def test_generate_raises_without_init() -> None:
224 """Test that generate raises when model has not been initialized."""
225 model = Wan22VideoGeneration()
226 # NOTE: intentionally NOT calling model.init() here
228 img = Image.new("RGB", (40, 30))
230 with pytest.raises(ValueError, match="Model not initialized"):
231 await model.generate(
232 img=img,
233 prompt="test prompt",
234 audio_path="/tmp/fake_audio.wav",
235 )
237 del model
238 gc.collect()
241def test_assert_model_init_wan_s2v_none() -> None:
242 """Test _assert_model_init raises when wan_s2v is None after init."""
243 model = Wan22VideoGeneration()
244 model.init()
245 model.wan_s2v = None
246 with pytest.raises(ValueError, match="WanS2V model not initialized"):
247 model._assert_model_init()
250def test_model_compile_no_op_no_torch_compile() -> None:
251 """Test model_compile is no-op when torch_compile=False."""
252 model = Wan22VideoGeneration()
253 model.init()
254 model.torch_compile = False
255 model.model_compile() # should not raise
258def test_model_compile_no_op_no_wan_s2v() -> None:
259 """Test model_compile is no-op when wan_s2v is None."""
260 model = Wan22VideoGeneration()
261 model.init()
262 model.torch_compile = True
263 model.wan_s2v = None
264 model.model_compile() # should not raise (wan_s2v is None guard)
267def test_model_compile_with_torch_compile() -> None:
268 """Test model_compile calls torch.compile when torch_compile=True and wan_s2v is set."""
269 model = Wan22VideoGeneration()
270 model.init()
271 model.torch_compile = True
272 assert model.wan_s2v is not None
273 model.model_compile() # torch.compile is mocked, should not raise
276@pytest.mark.asyncio
277async def test_get_rest_args_img_not_string() -> None:
278 """Test get_rest_args raises when img is not a string."""
279 model = Wan22VideoGeneration()
280 model.init()
282 with pytest.raises(ValueError, match="'img' parameter must be a base64-encoded string"):
283 await model.get_rest_args({"img": 12345})
286@pytest.mark.asyncio
287async def test_get_rest_args_audio_not_string() -> None:
288 """Test get_rest_args raises when audio is not a string."""
289 model = Wan22VideoGeneration()
290 model.init()
292 img = Image.new("RGB", (40, 30))
293 img_base64 = img_to_base64(img)
295 with pytest.raises(ValueError, match="'audio' parameter must be a base64-encoded string"):
296 await model.get_rest_args({
297 "img": img_base64,
298 "prompt": "test",
299 "audio": 99999,
300 })
303@pytest.mark.asyncio
304async def test_get_rest_args_tts_prompt_audio_not_string() -> None:
305 """Test get_rest_args raises when tts_prompt_audio is not a string."""
306 model = Wan22VideoGeneration()
307 model.init()
309 img = Image.new("RGB", (40, 30))
310 img_base64 = img_to_base64(img)
312 with pytest.raises(ValueError, match="'tts_prompt_audio' must be a base64-encoded string"):
313 await model.get_rest_args({
314 "img": img_base64,
315 "prompt": "test",
316 "enable_tts": True,
317 "tts_prompt_audio": 99999,
318 })
321@pytest.mark.asyncio
322async def test_get_rest_args_with_job_id() -> None:
323 """Test get_rest_args passes a deterministic audio destination when job_id is given."""
324 model = Wan22VideoGeneration()
325 model.init()
327 img = Image.new("RGB", (40, 30))
328 img_base64 = img_to_base64(img)
330 mock_media_utils.base64_to_audio_file.reset_mock()
332 await model.get_rest_args({
333 "img": img_base64,
334 "prompt": "test",
335 "audio": "dGVzdA==",
336 "job_id": "myjob123",
337 })
339 # Verify that the audio file helper was called with the expected destination path
340 mock_media_utils.base64_to_audio_file.assert_called_once_with(
341 "dGVzdA==", audio_path="/tmp/myjob123.wav"
342 )