Coverage for tests/test_run_benchmark.py: 100%

1#!/usr/bin/env python3 

2 

3import sys 

4import pytest 

5 

6from unittest.mock import patch 

7from unittest.mock import MagicMock 

8 

9from tests.torch_mock import TorchMock 

10from tests.diffusers_mock import DiffusersMock 

11 

12mock_torch = TorchMock() 

13mock_diffusers = DiffusersMock() 

14 

15sys.path.append("wrapper") 

16sys.path.append("wrapper/flux") 

17sys.path.append("wrapper/wan") 

18 

19mock_modules = { 

20 'torch.profiler': MagicMock(), 

21 'torch.hub': MagicMock(), 

22 'torch.version': MagicMock(), 

23 'torchvision': MagicMock(), 

24 'torch.distributed': MagicMock(), 

25 'torch.amp': MagicMock(), 

26 'torchvision.transforms': MagicMock(), 

27 'torchvision.transforms.functional': MagicMock(), 

28 'xfuser': MagicMock(), 

29 'xfuser.config': MagicMock(), 

30 'xfuser.core': MagicMock(), 

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

32 'xfuser.model_executor': MagicMock(), 

33 'xfuser.model_executor.models': MagicMock(), 

34 'xfuser.model_executor.models.transformers.transformer_flux': MagicMock(), 

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

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

37 'wan.modules': MagicMock(), 

38 'wan.modules.t5': MagicMock(), 

39 'wan.modules.clip': MagicMock(), 

40 'wan.modules.vae': MagicMock(), 

41 'wan.modules.model': MagicMock(), 

42 'wan.utils': MagicMock(), 

43 'wan.utils.utils': MagicMock(), 

44 'wan.utils.fm_solvers_unipc': MagicMock(), 

45 'wan.distributed': MagicMock(), 

46 'wan.distributed.fsdp': MagicMock(), 

47 'wan.distributed.xdit_context_parallel': MagicMock(), 

48} 

49mock_modules.update(mock_torch.get_sub_modules()) 

50mock_modules.update(mock_diffusers.get_sub_modules()) 

51 

52with patch.dict(sys.modules, mock_modules): 

53 from flux.run_flux_benchmark import main as run_flux_benchmark_main 

54 from wan.run_wan_benchmark import main as run_wan_benchmark_main 

55 from wan.run_wan_benchmark_batching import main as run_wan_benchmark_batching_main 

56 # import wan.run_wan_vae_decoder_benchmark 

57 

58 

59def test_run_flux_benchmark() -> None: 

60 with pytest.raises(ValueError): # engine_config, input_config = engine_args.create_config() 

61 run_flux_benchmark_main() 

62 

63 

64def test_run_wan_benchmark() -> None: 

65 with pytest.raises(FileNotFoundError): # No such file or directory: 'generated_image.png' 

66 run_wan_benchmark_main() 

67 

68 

69def test_run_wan_benchmark_batching_main() -> None: 

70 with pytest.raises(FileNotFoundError): # No such file or directory: 'generated_image.png' 

71 run_wan_benchmark_batching_main()