Coverage for tests/test_wrapper_realesrgan.py: 100%

1#!/usr/bin/env python3 

2 

3import sys 

4import pytest 

5 

6from unittest.mock import patch 

7from unittest.mock import MagicMock 

8from tests.torch_mock import TorchMock 

9 

10from PIL import Image 

11 

12mock_torch = TorchMock() 

13 

14sys.path.append("realesrgan") 

15 

16with patch.dict(sys.modules, { 

17 'nvidia_smi': MagicMock(), 

18 'imageio': MagicMock(), 

19 'cv2': MagicMock(), 

20 'RealESRGAN': MagicMock(), 

21 'torch': mock_torch, 

22}): 

23 from image_utils import img_to_base64 

24 from realesrgan.wrapper_realesrgan import RealESRGANGeneration 

25 

26 

27@pytest.mark.asyncio 

28async def test_wrapper_realesrgan() -> None: 

29 model = RealESRGANGeneration() 

30 assert model is not None 

31 assert model.model_name == "realesrgan" 

32 assert model.status == "initializing" 

33 assert len(model.models) == 0 

34 

35 with pytest.raises(ValueError): 

36 await model.generate(image=None) 

37 

38 model.init() 

39 assert model.status == "ok" 

40 assert len(model.models) == 3 

41 

42 health = model.get_health() 

43 assert health is not None 

44 assert len(health) >= 5 

45 timestamps = model.get_timestamps() 

46 assert timestamps is not None 

47 assert len(timestamps) >= 7 

48 

49 with pytest.raises(ValueError): 

50 await model.get_rest_args(None) 

51 await model.get_rest_args({}) 

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

53 img_base64 = img_to_base64(img) 

54 await model.get_rest_args({ 

55 "job_id": "unittest", 

56 "img": img_base64, 

57 "width": 80, 

58 "height": 60, 

59 }) 

60 

61 await model.warmup() 

62 

63 image_resized = await model.generate( 

64 image=img, 

65 width=160, 

66 height=120, 

67 output_type="pil" 

68 ) 

69 assert image_resized is not None 

70 

71 del model 

72 

73 

74def test_get_model_scaling_factor() -> None: 

75 """Test that get_model_scaling_factor picks the right (next-larger) scale.""" 

76 model = RealESRGANGeneration() 

77 model.init() 

78 

79 # 2x — exact match 

80 sf = model.get_model_scaling_factor(100, 100, 200, 200) 

81 assert sf == 2 

82 

83 # 3x ratio → should round up to 4x 

84 sf = model.get_model_scaling_factor(100, 100, 300, 300) 

85 assert sf == 4 

86 

87 # 4x — exact match 

88 sf = model.get_model_scaling_factor(100, 100, 400, 400) 

89 assert sf == 4 

90 

91 # 5x ratio → should round up to 8x (largest available) 

92 sf = model.get_model_scaling_factor(100, 100, 500, 500) 

93 assert sf == 8 

94 

95 # Asymmetric: width doubles, height quadruples → max factor 4 

96 sf = model.get_model_scaling_factor(100, 100, 200, 400) 

97 assert sf == 4 

98 

99 # Output smaller than input raises ValueError 

100 with pytest.raises(ValueError, match="must be larger than input"): 

101 model.get_model_scaling_factor(200, 200, 100, 100) 

102 

103 # Scaling factor exceeds max (8x) 

104 with pytest.raises(ValueError, match="Scaling factor"): 

105 model.get_model_scaling_factor(10, 10, 1000, 1000) 

106 

107 del model 

108 

109 

110def test_chunk_list_image() -> None: 

111 """Test _chunk_list_image distributes frames across ranks correctly.""" 

112 from PIL import Image 

113 

114 model = RealESRGANGeneration() 

115 

116 # Single rank (world_size == 1): all frames returned unchanged 

117 model.world_size = 1 

118 model.rank = 0 

119 frames = [Image.new("RGB", (10, 10)) for _ in range(4)] 

120 result = model._chunk_list_image(frames) 

121 assert result == frames 

122 

123 # Multi-rank (world_size == 2): rank 0 gets even-indexed frames 

124 model.world_size = 2 

125 model.rank = 0 

126 result = model._chunk_list_image(frames) 

127 assert len(result) == 4 

128 assert result[0] is frames[0] # index 0 → rank 0 

129 assert result[1] is None # index 1 → rank 1 

130 assert result[2] is frames[2] # index 2 → rank 0 

131 assert result[3] is None # index 3 → rank 1 

132 

133 # Multi-rank (world_size == 2): rank 1 gets odd-indexed frames 

134 model.rank = 1 

135 result = model._chunk_list_image(frames) 

136 assert result[0] is None 

137 assert result[1] is frames[1] 

138 assert result[2] is None 

139 assert result[3] is frames[3] 

140 

141 # Empty list always returns empty list 

142 model.world_size = 2 

143 model.rank = 0 

144 result = model._chunk_list_image([]) 

145 assert result == [] 

146 

147 del model