Coverage for tests/test_wrapper_imageresize.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 

14with patch.dict(sys.modules, { 

15 'nvidia_smi': MagicMock(), 

16 'imageio': MagicMock(), 

17 'cv2': MagicMock(), 

18 'torch': mock_torch, 

19}): 

20 from image_utils import img_to_base64 

21 from media_utils import video_frames_to_base64 

22 from imageresize.wrapper_imageresize import ImageResize 

23 

24 

25@pytest.mark.asyncio 

26async def test_e2e() -> None: 

27 model = ImageResize() 

28 assert model is not None 

29 assert model.model_name == "imageresize" 

30 

31 model.init() 

32 

33 health = model.get_health() 

34 assert health is not None 

35 

36 await model.warmup() 

37 

38 image = await model.generate(image=Image.new('RGB', (100, 100))) 

39 assert image is not None 

40 

41 

42@pytest.mark.asyncio 

43async def test_get_health() -> None: 

44 model = ImageResize() 

45 assert model is not None 

46 assert model.model_name == "imageresize" 

47 assert model.status == "initializing" 

48 

49 model.init() 

50 health = model.get_health() 

51 assert health is not None 

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

53 timestamps = model.get_timestamps() 

54 assert timestamps is not None 

55 

56 

57@pytest.mark.asyncio 

58async def test_get_rest_args() -> None: 

59 model = ImageResize() 

60 

61 with pytest.raises(ValueError): 

62 await model.get_rest_args(None) 

63 rest_args = await model.get_rest_args({}) 

64 assert rest_args is not None 

65 assert len(rest_args) == 2 

66 assert rest_args["task"] == "imageresize" 

67 rest_args = await model.get_rest_args({ 

68 "img": img_to_base64(Image.new('RGB', (100, 100))), 

69 }) 

70 assert len(rest_args) == 2 

71 rest_args = await model.get_rest_args({ 

72 "video": video_frames_to_base64([ 

73 Image.new('RGB', (100, 100)), 

74 Image.new('RGB', (100, 100)), 

75 ]), 

76 }) 

77 assert len(rest_args) == 2 

78 

79 

80@pytest.mark.asyncio 

81async def test_generate() -> None: 

82 model = ImageResize() 

83 

84 with pytest.raises(TypeError): 

85 await model.generate() 

86 image_resize = await model.generate( 

87 image=Image.new('RGB', (100, 100)), 

88 width=200, 

89 height=200) 

90 assert image_resize is not None 

91 assert image_resize.size == (200, 200) 

92 

93 video_resize = await model.generate( 

94 image=None, 

95 video=[ 

96 Image.new('RGB', (50, 25)), 

97 Image.new('RGB', (50, 25)), 

98 ], 

99 width=100, 

100 height=50) 

101 assert video_resize is not None 

102 assert len(video_resize) == 2 

103 assert video_resize[0].size == (100, 50) 

104 

105 with pytest.raises(ValueError): 

106 await model.generate(image=None)