Coverage for tests/test_wrapper_4kagent.py: 100%

1#!/usr/bin/env python3 

2 

3import os 

4import sys 

5import pytest 

6 

7from unittest.mock import AsyncMock 

8from unittest.mock import MagicMock 

9from unittest.mock import patch 

10from tests.torch_mock import TorchMock 

11 

12from PIL import Image 

13 

14mock_torch = TorchMock() 

15 

16sys.path.append("wrapper") 

17sys.path.append("wrapper/4kagent") 

18 

19# Only mock modules that are not available in the test environment 

20# (torch is replaced with a lightweight mock; nvidia_smi and colorlog 

21# need stubs since we have no GPU / colourised logging in CI). 

22# PyYAML and Pillow are real packages available in the test environment 

23# and must NOT be mocked so that file-I/O helpers work correctly. 

24mock_modules = { 

25 "torch": mock_torch, 

26 "nvidia_smi": MagicMock(), 

27 "colorlog": MagicMock(), 

28} 

29mock_modules.update(mock_torch.get_sub_modules()) 

30 

31with patch.dict(sys.modules, mock_modules): 

32 from image_utils import img_to_base64 

33 from wrapper_4kagent import Upscale4KAgent 

34 

35 

36@pytest.mark.asyncio 

37async def test_wrapper_4kagent_init() -> None: 

38 """Test that Upscale4KAgent initialises with correct defaults.""" 

39 model = Upscale4KAgent() 

40 assert model is not None 

41 assert model.model_name == "4kagent" 

42 assert model.status == "initializing" 

43 assert model.fourk_agent_dir is None 

44 del model 

45 

46 

47@pytest.mark.asyncio 

48async def test_wrapper_4kagent_generate_before_init() -> None: 

49 """generate() must raise when the model has not been initialised.""" 

50 model = Upscale4KAgent() 

51 img = Image.new("RGB", (64, 48)) 

52 with pytest.raises(ValueError, match="Model not initialized"): 

53 await model.generate(image=img) 

54 del model 

55 

56 

57@pytest.mark.asyncio 

58async def test_wrapper_4kagent_init_ok() -> None: 

59 """load_model() succeeds when FOURK_AGENT_DIR exists.""" 

60 model = Upscale4KAgent() 

61 

62 with patch("os.path.isdir", return_value=True), \ 

63 patch.object(model, "_write_config"): 

64 model.init() 

65 

66 assert model.status == "ok" 

67 assert model.fourk_agent_dir == Upscale4KAgent.FOURK_AGENT_DIR 

68 del model 

69 

70 

71@pytest.mark.asyncio 

72async def test_wrapper_4kagent_init_missing_dir() -> None: 

73 """load_model() raises RuntimeError when FOURK_AGENT_DIR is absent.""" 

74 model = Upscale4KAgent() 

75 with patch("os.path.isdir", return_value=False): 

76 with pytest.raises(RuntimeError, match="4KAgent not found"): 

77 model.init() 

78 del model 

79 

80 

81@pytest.mark.asyncio 

82async def test_wrapper_4kagent_get_rest_args_missing_json() -> None: 

83 """get_rest_args() raises on None or non-dict input.""" 

84 model = Upscale4KAgent() 

85 with pytest.raises(ValueError, match="Missing JSON body"): 

86 await model.get_rest_args(None) # type: ignore[arg-type] 

87 with pytest.raises(ValueError, match="Missing JSON body"): 

88 await model.get_rest_args("not a dict") # type: ignore[arg-type] 

89 del model 

90 

91 

92@pytest.mark.asyncio 

93async def test_wrapper_4kagent_get_rest_args_missing_img() -> None: 

94 """get_rest_args() raises when the 'img' key is absent.""" 

95 model = Upscale4KAgent() 

96 with pytest.raises(ValueError, match="Missing 'img' parameter"): 

97 await model.get_rest_args({}) 

98 del model 

99 

100 

101@pytest.mark.asyncio 

102async def test_wrapper_4kagent_get_rest_args_success() -> None: 

103 """get_rest_args() returns the correct structure for a valid payload.""" 

104 model = Upscale4KAgent() 

105 

106 img = Image.new("RGB", (40, 30)) 

107 img_base64 = img_to_base64(img) 

108 assert img_base64 is not None 

109 

110 result = await model.get_rest_args({ 

111 "job_id": "test-job", 

112 "img": img_base64, 

113 "profile_name": "ExpSR_s4_F", 

114 "tool_run_gpu_id": 1, 

115 }) 

116 

117 assert result["task"] == "4kagent" 

118 args = result["args"] 

119 assert args["job_id"] == "test-job" 

120 assert args["profile_name"] == "ExpSR_s4_F" 

121 assert args["tool_run_gpu_id"] == 1 

122 assert args["image"] is not None 

123 del model 

124 

125 

126@pytest.mark.asyncio 

127async def test_wrapper_4kagent_get_rest_args_defaults() -> None: 

128 """get_rest_args() applies default profile_name and tool_run_gpu_id.""" 

129 model = Upscale4KAgent() 

130 img = Image.new("RGB", (40, 30)) 

131 img_base64 = img_to_base64(img) 

132 assert img_base64 is not None 

133 

134 result = await model.get_rest_args({"img": img_base64}) 

135 

136 assert result["args"]["profile_name"] == Upscale4KAgent.DEFAULT_PROFILE 

137 assert result["args"]["tool_run_gpu_id"] == 0 

138 del model 

139 

140 

141@pytest.mark.asyncio 

142async def test_wrapper_4kagent_generate_no_image() -> None: 

143 """generate() raises ValueError when no image is provided.""" 

144 model = Upscale4KAgent() 

145 with patch("os.path.isdir", return_value=True), \ 

146 patch.object(model, "_write_config"): 

147 model.init() 

148 

149 with pytest.raises(ValueError, match="input image is required"): 

150 await model.generate(image=None) 

151 del model 

152 

153 

154@pytest.mark.asyncio 

155async def test_wrapper_4kagent_generate_success() -> None: 

156 """generate() calls _run_4kagent and returns the PIL result.""" 

157 model = Upscale4KAgent() 

158 with patch("os.path.isdir", return_value=True), \ 

159 patch.object(model, "_write_config"): 

160 model.init() 

161 

162 expected_image = Image.new("RGB", (1024, 1024), color=(200, 200, 200)) 

163 input_image = Image.new("RGB", (64, 48)) 

164 

165 with patch.object(model, "_run_4kagent"), \ 

166 patch.object(model, "_load_result", return_value=expected_image): 

167 result = await model.generate( 

168 image=input_image, 

169 profile_name="ExpSR_s4_P", 

170 tool_run_gpu_id=0, 

171 job_id="unittest", 

172 ) 

173 

174 assert result is expected_image 

175 assert not model.running # Must be cleared after completion 

176 del model 

177 

178 

179@pytest.mark.asyncio 

180async def test_wrapper_4kagent_generate_error() -> None: 

181 """generate() propagates failure from _run_4kagent as RuntimeError.""" 

182 model = Upscale4KAgent() 

183 with patch("os.path.isdir", return_value=True), \ 

184 patch.object(model, "_write_config"): 

185 model.init() 

186 

187 def fail_run(*args: object, **kwargs: object) -> None: 

188 raise RuntimeError("4KAgent failed: error") 

189 

190 input_image = Image.new("RGB", (64, 48)) 

191 with patch.object(model, "_run_4kagent", side_effect=fail_run): 

192 with pytest.raises(RuntimeError, match="4KAgent failed"): 

193 await model.generate(image=input_image) 

194 

195 assert not model.running 

196 del model 

197 

198 

199@pytest.mark.asyncio 

200async def test_wrapper_4kagent_health() -> None: 

201 """get_health() returns required keys including 4KAgent-specific ones.""" 

202 model = Upscale4KAgent() 

203 with patch("os.path.isdir", return_value=True), \ 

204 patch.object(model, "_write_config"): 

205 model.init() 

206 

207 health = model.get_health() 

208 assert "model_name" in health 

209 assert "status" in health 

210 assert "fourk_agent_dir" in health 

211 del model 

212 

213 

214@pytest.mark.asyncio 

215async def test_wrapper_4kagent_warmup() -> None: 

216 """warmup() delegates to generate() with a synthetic image.""" 

217 model = Upscale4KAgent() 

218 with patch("os.path.isdir", return_value=True), \ 

219 patch.object(model, "_write_config"): 

220 model.init() 

221 

222 warmup_image = Image.new("RGB", (1024, 1024)) 

223 with patch.object(model, "generate", new=AsyncMock(return_value=warmup_image)) as mock_gen: 

224 await model.warmup() 

225 mock_gen.assert_called_once() 

226 # The warmup call must pass an image argument 

227 call_kwargs = mock_gen.call_args 

228 assert call_kwargs.kwargs.get("image") is not None or call_kwargs.args 

229 del model 

230 

231 

232def test_wrapper_4kagent_write_config() -> None: 

233 """_write_config() writes a YAML file populated from environment variables.""" 

234 import yaml 

235 import tempfile 

236 

237 model = Upscale4KAgent() 

238 

239 with tempfile.TemporaryDirectory() as tmpdir: 

240 model.fourk_agent_dir = tmpdir 

241 

242 env = { 

243 "LLAMA_API_KEY": "test-llama-key", 

244 "OPENAI_API_KEY": "test-openai-key", 

245 "AZURE_OPENAI_API_KEY": "", 

246 "AZURE_OPENAI_ENDPOINT": "", 

247 "AZURE_OPENAI_MODEL": "", 

248 "AZURE_OPENAI_API_VERSION": "", 

249 } 

250 with patch.dict(os.environ, env, clear=False): 

251 model._write_config() 

252 

253 config_path = os.path.join(tmpdir, "config.yml") 

254 assert os.path.isfile(config_path) 

255 

256 with open(config_path) as fh: 

257 config = yaml.safe_load(fh) 

258 

259 assert config["LLAMA"]["API_KEY"] == "test-llama-key" 

260 assert config["GPT"]["API_KEY"] == "test-openai-key" 

261 assert "AZUREGPT" in config 

262 

263 del model 

264 

265 

266def test_load_result_no_results() -> None: 

267 """_load_result() raises ValueError when no result.png exists.""" 

268 import tempfile 

269 

270 model = Upscale4KAgent() 

271 with tempfile.TemporaryDirectory() as tmpdir: 

272 with pytest.raises(ValueError, match="No result.png found"): 

273 model._load_result(tmpdir) 

274 del model 

275 

276 

277def test_load_result_finds_image() -> None: 

278 """_load_result() returns a PIL Image from the expected output structure.""" 

279 import tempfile 

280 

281 model = Upscale4KAgent() 

282 

283 with tempfile.TemporaryDirectory() as tmpdir: 

284 # Simulate 4KAgent output: <output_dir>/input/<step>/result.png 

285 result_dir = os.path.join(tmpdir, "input", "step_001") 

286 os.makedirs(result_dir) 

287 result_path = os.path.join(result_dir, "result.png") 

288 Image.new("RGB", (1024, 768), color=(10, 20, 30)).save(result_path) 

289 

290 result = model._load_result(tmpdir) 

291 

292 assert isinstance(result, Image.Image) 

293 assert result.size == (1024, 768) 

294 del model