Coverage for tests/test_wrapper_qwenimage.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 qwenimage.wrapper_qwenimage import QwenImageGeneration 

28 

29 

30@pytest.mark.asyncio 

31async def test_wrapper_qwenimage() -> None: 

32 model = QwenImageGeneration() 

33 assert model is not None 

34 assert model.model_name == "qwenimage" 

35 assert model.status == "initializing" 

36 

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

38 await model.generate(64, 48, "test prompt") 

39 

40 model.init() 

41 assert model.status == "ok" 

42 

43 # Mock pipeline return object 

44 mock_output = MagicMock() 

45 mock_output.images = [ 

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

47 ] 

48 model.pipeline = MagicMock(return_value=mock_output) 

49 model.pipeline.vae_scale_factor = 8 

50 

51 health = model.get_health() 

52 assert health is not None 

53 assert health["model_name"] == "qwenimage" 

54 assert health["running"] is False 

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

56 assert "load_timer" in health 

57 assert "gen_timer" in health 

58 

59 timestamps = model.get_timestamps() 

60 assert timestamps is not None 

61 

62 with pytest.raises(ValueError): 

63 await model.get_rest_args(None) 

64 with pytest.raises(ValueError): 

65 await model.get_rest_args({}) 

66 await model.get_rest_args({ 

67 "job_id": "unittest", 

68 "prompt": "Test prompt", 

69 "width": 80, 

70 "height": 60, 

71 "seed": 7, 

72 }) 

73 

74 await model.warmup() 

75 

76 image = await model.generate( 

77 prompt="Test prompt", 

78 height=1024, 

79 width=1024) 

80 assert image is not None 

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

82 

83 del model 

84 

85 

86@pytest.mark.asyncio 

87async def test_wrapper_qwenimage_assert_args() -> None: 

88 """_assert_args raises for image sizes not evenly divisible across GPUs.""" 

89 model = QwenImageGeneration() 

90 model.init() 

91 model.pipeline = MagicMock(return_value=MagicMock(images=[ 

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

93 ])) 

94 model.pipeline.vae_scale_factor = 8 # latent factor = 8 

95 

96 # Single GPU (world_size=1): any size accepted 

97 model.world_size = 1 

98 model.rank = 0 

99 image = await model.generate(prompt="test", height=512, width=512) 

100 assert image is not None 

101 

102 # Multi-GPU (world_size=3): latent shape 1024 is not divisible by 3 → raises 

103 model.world_size = 3 

104 with pytest.raises(ValueError, match="not supported for"): 

105 await model.generate(prompt="test", height=512, width=512) 

106 

107 # Multi-GPU (world_size=4): latent shape 1024 is divisible by 4 → OK 

108 model.world_size = 4 

109 image = await model.generate(prompt="test", height=512, width=512) 

110 assert image is not None 

111 

112 del model