Coverage for tests/streamwise_app/test_lmm_generator.py: 99%
151 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
2"""
3Unit tests for LMM Generator.
4"""
6import os
7import sys
8import pytest
10from PIL import Image
12from unittest.mock import patch
13from unittest.mock import MagicMock
15from pydantic import ValidationError
17# Add current path
18sys.path.append(os.getcwd())
20from tests.test_utils import temp_sys_path
22from tests.torch_mock import TorchMock
23from tests.k8s_mock import K8sMock
24from tests.fantasytalking_mock import FantasyTalkingMock
26mock_torch = TorchMock()
27mock_k8s = K8sMock()
28mock_ft = FantasyTalkingMock()
30mock_modules = {
31 "imageio": MagicMock(),
32 "tabulate": MagicMock(),
33 "soundfile": MagicMock(),
34}
35mock_modules.update(mock_torch.get_sub_modules())
36mock_modules.update(mock_k8s.get_sub_modules())
37mock_modules.update(mock_ft.get_sub_modules())
39with patch.dict(sys.modules, mock_modules):
40 from file_utils import read_file_base64
42 with temp_sys_path("apps"):
43 from apps.lmm_generator import LMMGenerator
45 from apps.client import ServiceRequest
46 from apps.client import RequestStatus
49@pytest.mark.asyncio
50async def test_gen_image() -> None:
51 service_manager = MagicMock()
52 job_id = "gen_image"
53 gen = LMMGenerator("streamwise_app", job_id, service_manager)
54 try:
55 # with pytest.raises(ServiceError):
56 with pytest.raises(Exception, match="Service 'flux' not found"):
57 await gen.gen_image(
58 prompt="Test prompt",
59 task_id="test_gen_image")
60 finally:
61 await gen.stop()
62 del gen
63 del service_manager
66@pytest.mark.asyncio
67async def test_gen_edit_image() -> None:
68 service_manager = MagicMock()
69 job_id = "gen_edit_image"
70 gen = LMMGenerator("streamwise_app", job_id, service_manager)
71 try:
72 # with pytest.raises(ServiceError):
73 with pytest.raises(Exception, match="Service 'fluxkontext' not found"):
74 await gen.gen_edit_image(
75 prompt="Test prompt",
76 image=Image.new('RGB', (100, 100)),
77 task_id="test_gen_edit_image")
78 finally:
79 await gen.stop()
80 del gen
81 del service_manager
84@pytest.mark.asyncio
85async def test_gen_extract_characters() -> None:
86 service_manager = MagicMock()
87 job_id = "job_test_gen_extract_characters"
88 gen = LMMGenerator("streamwise_app", job_id, service_manager)
89 try:
90 # with pytest.raises(ServiceError):
91 with pytest.raises(Exception, match="Service 'yolo' not found"):
92 await gen.gen_extract_characters(
93 image=Image.new('RGB', (100, 100)),
94 task_id="gen_extract_characters")
95 finally:
96 await gen.stop()
97 del gen
98 del service_manager
101@pytest.mark.asyncio
102async def test_gen_video() -> None:
103 service_manager = MagicMock()
104 job_id = "job_test_gen_video"
105 gen = LMMGenerator("streamwise_app", job_id, service_manager)
106 try:
107 # with pytest.raises(ServiceError):
108 with pytest.raises(Exception, match="Service 'hunyuanframepackf1' not found"):
109 await gen.gen_video(
110 prompt="Test prompt",
111 img=Image.new('RGB', (100, 100)),
112 task_id="test_gen_video")
113 finally:
114 await gen.stop()
115 del gen
116 del service_manager
119@pytest.mark.asyncio
120async def test_gen_video_audio_from_img() -> None:
121 service_manager = MagicMock()
122 job_id = "job_test_gen_video_audio_from_img"
123 gen = LMMGenerator("streamwise_app", job_id, service_manager)
124 try:
125 audio_base64 = await read_file_base64("tests/data/audio_4675.wav")
126 # with pytest.raises(ServiceError):
127 with pytest.raises(Exception, match="Service 'fantasytalking' not found"):
128 await gen.gen_video_audio_from_img(
129 prompt="Test prompt",
130 audio_base64=audio_base64,
131 img=Image.new('RGB', (100, 100)),
132 task_id="test")
133 finally:
134 await gen.stop()
135 del gen
136 del service_manager
139@pytest.mark.asyncio
140async def test_gen_video_audio_from_video() -> None:
141 service_manager = MagicMock()
142 job_id = "job_test_gen_video_audio_from_video"
143 gen = LMMGenerator("streamwise_app", job_id, service_manager)
144 try:
145 audio_base64 = await read_file_base64("tests/data/audio_4675.wav")
146 # with pytest.raises(ServiceError):
147 with pytest.raises(Exception, match="Service 'fantasytalking' not found"):
148 await gen.gen_video_audio_from_video(
149 prompt="Test prompt",
150 audio_base64=audio_base64,
151 video=[
152 Image.new('RGB', (100, 100)),
153 Image.new('RGB', (100, 100)),
154 Image.new('RGB', (100, 100)),
155 ],
156 task_id="test")
157 finally:
158 await gen.stop()
159 del gen
160 del service_manager
163@pytest.mark.asyncio
164async def test_get_video_last_latents() -> None:
165 service_manager = MagicMock()
166 job_id = "job_test_get_video_last_latents"
167 gen = LMMGenerator("streamwise_app", job_id, service_manager)
168 try:
169 index, latents = await gen.get_video_last_latents(
170 base_url="http://test:8080",
171 task_id="test_get_video_last_latents")
172 assert index == -1
173 assert latents is None
174 finally:
175 await gen.stop()
176 del gen
177 del service_manager
180@pytest.mark.asyncio
181async def test_gen_video_from_latents() -> None:
182 service_manager = MagicMock()
183 job_id = "job_test_gen_video_from_latents"
184 gen = LMMGenerator("streamwise_app", job_id, service_manager)
185 try:
186 with pytest.raises(ValidationError):
187 await gen.gen_video_from_latents(
188 latents=mock_torch.randn(1, 4, 8, 8),
189 task_id="test")
190 finally:
191 await gen.stop()
192 del gen
193 del service_manager
196@pytest.mark.asyncio
197async def test_gen_intermediate_video_frames() -> None:
198 job_id = "gen_intermediate_video_frames"
199 service_manager = MagicMock()
200 gen = LMMGenerator("streamwise_app", job_id, service_manager)
201 try:
202 request = ServiceRequest(
203 request_id="request_gen_video",
204 service_name="video-gen",
205 payload_bytes=b"test")
206 request.status = RequestStatus.COMPLETED
208 async for frame in gen.gen_intermediate_video_frames(
209 base_url="http://test:8080",
210 video_gen_request=request,
211 task_id="test_gen_intermediate_video_frames"
212 ):
213 assert frame is not None
214 finally:
215 await gen.stop()
216 del gen
217 del service_manager
220@pytest.mark.asyncio
221async def test_gen_image_upscale() -> None:
222 service_manager = MagicMock()
223 job_id = "gen_image_upscale"
224 gen = LMMGenerator("streamwise_app", job_id, service_manager)
225 try:
226 # with pytest.raises(ServiceError):
227 with pytest.raises(Exception, match="Service 'realesrgan' not found"):
228 await gen.gen_image_upscale(
229 image=Image.new('RGB', (40, 40)),
230 task_id="test_gen_image_upscale")
231 finally:
232 await gen.stop()
233 del gen
234 del service_manager
237@pytest.mark.asyncio
238async def test_gen_video_upscale() -> None:
239 service_manager = MagicMock()
240 job_id = "gen_video_upscale"
241 gen = LMMGenerator("streamcast", job_id, service_manager)
242 try:
243 # with pytest.raises(ServiceError):
244 with pytest.raises(Exception, match="Service 'realesrgan' not found"):
245 await gen.gen_video_upscale(
246 video_binary=b"AAAAAAAA",
247 task_id="test_gen_video_upscale")
248 finally:
249 await gen.stop()
250 del gen
251 del service_manager