Coverage for tests/test_wrapper_kokoro.py: 100%

1#!/usr/bin/env python3 

2 

3import sys 

4import gc 

5import pytest 

6 

7from unittest.mock import patch 

8from unittest.mock import MagicMock 

9from tests.torch_mock import TorchMock 

10 

11mock_torch = TorchMock() 

12 

13sys.path.append("kokoro") 

14 

15mock_modules = { 

16 'kokoro.KPipeline': MagicMock(), 

17 'nvidia_smi': MagicMock(), 

18 'imageio': MagicMock(), 

19 'cv2': MagicMock(), 

20 'torch': mock_torch, 

21} 

22mock_modules.update(mock_torch.get_sub_modules()) 

23 

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

25 from kokoro.wrapper_kokoro import KokoroGeneration 

26 from kokoro.wrapper_kokoro import Language 

27 

28 

29@pytest.mark.asyncio 

30async def test_basic() -> None: 

31 model = KokoroGeneration() 

32 assert model is not None 

33 assert model.model_name == "kokoro" 

34 assert model.status == "initializing" 

35 

36 model.init() 

37 assert model.status == "ok" 

38 health = model.get_health() 

39 assert health is not None 

40 assert len(health) >= 5 

41 timestamps = model.get_timestamps() 

42 assert timestamps is not None 

43 assert len(timestamps) >= 4 

44 

45 with pytest.raises(ValueError): 

46 await model.get_rest_args(None) 

47 with pytest.raises(ValueError): 

48 await model.get_rest_args({}) 

49 await model.get_rest_args({ 

50 "text": "Test text" 

51 }) 

52 

53 await model.warmup() 

54 

55 audio_path = await model.generate( 

56 text="Test text", 

57 output_type="audio_path") 

58 # TODO make it not None 

59 assert audio_path is None 

60 

61 del model 

62 gc.collect() 

63 

64 

65def test_language_enum() -> None: 

66 """All Language enum members should have their correct string codes.""" 

67 assert Language.AMERICAN_ENGLISH.value == "a" 

68 assert Language.BRITISH_ENGLISH.value == "b" 

69 assert Language.SPANISH.value == "e" 

70 assert Language.FRENCH.value == "f" 

71 assert Language.HINDI.value == "h" 

72 assert Language.ITALIAN.value == "i" 

73 assert Language.BRAZILIAN_PORTUGUESE.value == "p" 

74 assert Language.JAPANESE.value == "j" 

75 assert Language.MANDARIN_CHINESE.value == "z" 

76 

77 # Language is a str enum — its value is equal to the plain string 

78 assert Language.AMERICAN_ENGLISH == "a" 

79 

80 

81@pytest.mark.asyncio 

82async def test_get_rest_args_optional_params() -> None: 

83 """get_rest_args should return defaults and accept custom voice/speed/lang.""" 

84 model = KokoroGeneration() 

85 model.init() 

86 

87 # Defaults: voice=af_heart, speed=1.0, lang_code=a (American English) 

88 args = await model.get_rest_args({"text": "hello"}) 

89 assert args["task"] == "kokoro" 

90 inner = args["args"] 

91 assert inner["text"] == "hello" 

92 assert inner["voice"] == "af_heart" 

93 assert inner["speed"] == 1.0 

94 assert inner["lang_code"] == "a" 

95 assert inner["job_id"] is None 

96 

97 # Custom voice, speed, and language 

98 args = await model.get_rest_args({ 

99 "job_id": "j1", 

100 "text": "bonjour", 

101 "voice": "bf_emma", 

102 "speed": "0.8", 

103 "lang_code": "f", 

104 }) 

105 inner = args["args"] 

106 assert inner["job_id"] == "j1" 

107 assert inner["text"] == "bonjour" 

108 assert inner["voice"] == "bf_emma" 

109 assert inner["speed"] == pytest.approx(0.8) 

110 assert inner["lang_code"] == "f" 

111 

112 # Pass a Language enum value directly (should be treated as its string value) 

113 args = await model.get_rest_args({ 

114 "text": "nihao", 

115 "lang_code": Language.MANDARIN_CHINESE, 

116 }) 

117 assert args["args"]["lang_code"] == Language.MANDARIN_CHINESE 

118 

119 del model 

120 gc.collect()