Coverage for tests/test_wrapper_qwenimageedit.py: 100%

1#!/usr/bin/env python3 

2 

3import sys 

4import pytest 

5 

6from PIL import Image 

7 

8from unittest.mock import patch 

9from unittest.mock import MagicMock 

10from tests.torch_mock import TorchMock 

11from tests.diffusers_mock import DiffusersMock 

12 

13mock_torch = TorchMock() 

14mock_diffusers = DiffusersMock() 

15 

16sys.path.append("wrapper") 

17 

18mock_modules = { 

19 "distvae.modules.adapters.vae.decoder_adapters": MagicMock(), 

20 "xfuser.config": MagicMock(), 

21 "xfuser.core.distributed.group_coordinator": MagicMock(), 

22} 

23mock_modules.update(mock_torch.get_sub_modules()) 

24mock_modules.update(mock_diffusers.get_sub_modules()) 

25 

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

27 from image_utils import img_to_base64 

28 from qwenimageedit.wrapper_qwenimageedit import QwenImageEditGeneration 

29 

30 

31@pytest.mark.asyncio 

32async def test_wrapper_qwenimage() -> None: 

33 model = QwenImageEditGeneration() 

34 assert model is not None 

35 assert model.model_name == "qwenimageedit" 

36 assert model.status == "initializing" 

37 

38 img_input = Image.new("RGB", (64, 48), color="red") 

39 img_base64 = img_to_base64(img_input) 

40 

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

42 await model.generate( 

43 img=img_input, 

44 width=64, 

45 height=48, 

46 prompt="test prompt") 

47 

48 model.init() 

49 assert model.status == "ok" 

50 

51 # Mock pipeline return object 

52 mock_output = MagicMock() 

53 mock_output.images = [ 

54 Image.new("RGB", (64, 48), color="red") 

55 ] 

56 model.pipeline = MagicMock(return_value=mock_output) 

57 model.pipeline.vae_scale_factor = 8 

58 

59 health = model.get_health() 

60 assert health is not None 

61 assert health["model_name"] == "qwenimageedit" 

62 assert health["running"] is False 

63 assert health["status"] == "ok" 

64 assert "load_timer" in health 

65 assert "gen_timer" in health 

66 

67 timestamps = model.get_timestamps() 

68 assert timestamps is not None 

69 

70 with pytest.raises(ValueError): 

71 await model.get_rest_args(None) 

72 with pytest.raises(ValueError): 

73 await model.get_rest_args({}) 

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

75 await model.get_rest_args({ 

76 "job_id": "unittest", 

77 "prompt": "Test prompt", 

78 "width": 80, 

79 "height": 60, 

80 "seed": 7, 

81 }) 

82 await model.get_rest_args({ 

83 "job_id": "unittest", 

84 "img": img_base64, 

85 "prompt": "Test prompt", 

86 "width": 80, 

87 "height": 60, 

88 "seed": 7, 

89 }) 

90 

91 await model.warmup() 

92 

93 image = await model.generate( 

94 img=img_input, 

95 prompt="Test prompt", 

96 height=1024, 

97 width=1024) 

98 assert image is not None 

99 assert image.size == (64, 48) # Returns the mock value 

100 

101 del model