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

10mock_torch = TorchMock() 

11 

12sys.path.append("wrapper") 

13with patch.dict(sys.modules, { 

14 'nvidia_smi': MagicMock(), 

15 'imageio': MagicMock(), 

16 'cv2': MagicMock(), 

17 'torch': mock_torch, 

18 'torch.amp': MagicMock(), 

19 'torch.distributed': MagicMock(), 

20 'xfuser': MagicMock(), 

21 'xfuser.config': MagicMock(), 

22 'xfuser.core': MagicMock(), 

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

24 'xfuser.model_executor': MagicMock(), 

25 'xfuser.model_executor.layers': MagicMock(), 

26 'xfuser.model_executor.layers.attention_processor': MagicMock(), 

27 'tokenizer.tokenizer_image.vq_model': MagicMock(), 

28 'language.t5': MagicMock(), 

29 'autoregressive.models.gpt': MagicMock(), 

30 'autoregressive.models.generate': MagicMock(), 

31}): 

32 from llamagen.wrapper_llamagen import LlamaGenGeneration 

33 

34 

35@pytest.mark.asyncio 

36async def test_wrapper_llamagen() -> None: 

37 model = LlamaGenGeneration() 

38 assert model is not None 

39 assert model.model_name == "llamagen" 

40 assert model.status == "initializing" 

41 

42 with pytest.raises(AssertionError): # vq_model not set 

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

44 

45 with pytest.raises(AttributeError, match="'LlamaGenGeneration' object has no attribute 'gpt_type'"): 

46 model.init() 

47 assert model.status == "failed" 

48 

49 # Mock pipeline return object 

50 mock_output = MagicMock() 

51 mock_output.images = ["image"] 

52 model.pipeline = MagicMock(return_value=mock_output) 

53 model.pipeline.vae_scale_factor = 8 

54 model.pipeline.gpt_type = "llama" 

55 model.vq_model_name = "vq_model" 

56 model.gpt_model_name = "gpt_model" 

57 model.t5_model_type = "t5_model" 

58 model.gpt_type = "llama" 

59 

60 with pytest.raises(FileNotFoundError, match="T5 model directory 'google' does not exist."): 

61 model.init() 

62 

63 health = model.get_health() 

64 assert health is not None 

65 assert "dtype" in health 

66 assert "gpu" in health 

67 assert "world_size" in health 

68 

69 timestamps = model.get_timestamps() 

70 assert timestamps is not None 

71 

72 with pytest.raises(ValueError): 

73 await model.get_rest_args(None) 

74 with pytest.raises(ValueError): 

75 await model.get_rest_args({}) 

76 await model.get_rest_args({ 

77 "job_id": "unittest", 

78 "prompt": "Test prompt", 

79 "width": 80, 

80 "height": 60, 

81 "seed": 7, 

82 }) 

83 

84 with pytest.raises(AssertionError): 

85 await model.warmup() 

86 

87 with pytest.raises(AssertionError): 

88 await model.generate(prompt="Test prompt") 

89 # assert image is not None 

90 

91 del model