Coverage for tests/test_wrapper_mock.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 mock.wrapper_mock import MockGeneration 

21 

22 

23@pytest.mark.asyncio 

24async def test_e2e() -> None: 

25 model = MockGeneration() 

26 assert model is not None 

27 assert model.model_name == "mock" 

28 assert model.status == "initializing" 

29 

30 model.init() 

31 

32 health = model.get_health() 

33 assert health is not None 

34 

35 await model.warmup() 

36 

37 image = await model.generate( 

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

39 output_type="pil" 

40 ) 

41 assert image is not None 

42 assert isinstance(image, Image.Image) 

43 

44 

45@pytest.mark.asyncio 

46async def test_generate() -> None: 

47 model = MockGeneration() 

48 

49 image = await model.generate( 

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

51 output_type="pil" 

52 ) 

53 assert image is not None 

54 assert isinstance(image, Image.Image) 

55 

56 jsonl = await model.generate(output_type="jsonl") 

57 assert jsonl is not None 

58 assert isinstance(jsonl, str) 

59 

60 tensor = await model.generate(output_type="tensor") 

61 assert tensor is not None 

62 assert isinstance(tensor, str) 

63 

64 

65# TODO fix this test 

66""" 

67@pytest.mark.asyncio 

68async def test_http_server() -> None: 

69 with patch.dict(sys.modules, { 

70 'nvidia_smi': MagicMock(), 

71 'torch': mock_torch, 

72 'torch.distributed': MagicMock(), 

73 'imageio': MagicMock(), 

74 'xfuser': MagicMock(), 

75 'xfuser.config': MagicMock(), 

76 'xfuser.core': MagicMock(), 

77 'xfuser.core.distributed': MagicMock(), 

78 }): 

79 from run_httpserver import main 

80 

81 test_args = ["run_http_server.py", "--mock"] 

82 with patch.object(sys, "argv", test_args): 

83 await main() 

84"""