Coverage for tests\unit\test_spinner.py: 100%

107 statements  

« prev     ^ index     » next       coverage.py v7.6.1, created at 2024-10-15 20:56 -0600

1from typing import Callable, Any 

2import time 

3 

4 

5import io 

6 

7from muutils.spinner import spinner_decorator, SpinnerContext, Spinner 

8 

9 

10def test_spinner_simple(): 

11 @spinner_decorator(update_interval=0.05) 

12 def long_running_function_simple() -> str: 

13 """ 

14 An example function decorated with spinner_decorator, using only the spinner and elapsed time. 

15 

16 Returns: 

17 str: A completion message. 

18 """ 

19 for _ in range(1): 

20 time.sleep(0.1) # Simulate some work 

21 return "Simple function completed" 

22 

23 print("\nRunning simple function with only spinner and elapsed time:") 

24 result2: str = long_running_function_simple() 

25 print(result2) 

26 

27 

28def test_spinner_complex(): 

29 # Example usage 

30 @spinner_decorator( 

31 message="Current value: ", 

32 mutable_kwarg_key="update_status", 

33 update_interval=0.05, 

34 ) 

35 def long_running_function_with_status( 

36 normal_arg: int, update_status: Callable[[Any], None] 

37 ) -> str: 

38 """ 

39 An example function decorated with spinner_decorator, using all features. 

40 

41 Args: 

42 normal_arg (int): A normal argument to demonstrate that other arguments still work. 

43 update_status (Callable[[Any], None]): Function to update the status displayed by the spinner. 

44 

45 Returns: 

46 str: A completion message. 

47 """ 

48 for i in range(normal_arg): 

49 time.sleep(0.1) # Simulate some work 

50 update_status(f"Step {i+1} of {normal_arg}") 

51 return "Function with status completed" 

52 

53 # Run the example functions 

54 print("Running function with status updates:") 

55 result1: str = long_running_function_with_status(1) 

56 print(result1) 

57 

58 

59def test_spinner_decorator_bare(): 

60 @spinner_decorator() 

61 def example_function(): 

62 return "Done" 

63 

64 result = example_function() 

65 assert result == "Done" 

66 

67 

68def test_spinner_ctx_mgr(): 

69 with SpinnerContext(message="Current value: ", update_interval=0.05) as spinner: 

70 for i in range(1): 

71 time.sleep(0.1) 

72 spinner.update_value(f"Step {i+1}") 

73 print("Done!") 

74 

75 

76def test_spinner_initialization(): 

77 spinner = Spinner() 

78 assert isinstance(spinner, Spinner) 

79 assert isinstance(spinner.spinner_chars, list) 

80 assert isinstance(spinner.update_interval, float) 

81 assert isinstance(spinner.spinner_complete, str) 

82 assert isinstance(spinner.current_value, str) 

83 assert isinstance(spinner.message, str) 

84 assert isinstance(spinner.format_string, str) 

85 assert hasattr(spinner.output_stream, "write") 

86 assert callable(spinner.output_stream.write) 

87 assert callable(spinner.update_value) 

88 

89 assert spinner.spinner_chars == ["|", "/", "-", "\\"] 

90 assert spinner.format_string == "\r{spinner} ({elapsed_time:.2f}s) {message}{value}" 

91 assert spinner.update_interval == 0.1 

92 

93 

94def test_spinner_update_value(): 

95 spinner = Spinner() 

96 spinner.update_value("Test") 

97 assert spinner.current_value == "Test" 

98 

99 

100def test_spinner_context_manager(): 

101 string_io = io.StringIO() 

102 with SpinnerContext(output_stream=string_io): 

103 pass 

104 assert string_io.getvalue().endswith("\n") 

105 

106 

107@spinner_decorator() 

108def example_function(): 

109 return "Done" 

110 

111 

112def test_spinner_decorator(): 

113 result = example_function() 

114 assert result == "Done" 

115 

116 

117def test_spinner_custom_chars(): 

118 spinner = Spinner(spinner_chars=["A", "B", "C"]) 

119 assert spinner.spinner_chars == ["A", "B", "C"] 

120 

121 

122def test_spinner_custom_time_format(): 

123 spinner = Spinner(format_string="[{elapsed_time:.1f}s]") 

124 assert spinner.format_string == "[{elapsed_time:.1f}s]" 

125 

126 

127def test_spinner_context_manager_with_updates(): 

128 string_io = io.StringIO() 

129 with SpinnerContext( 

130 message="Status: ", output_stream=string_io, update_interval=0.05 

131 ) as spinner: 

132 spinner.update_value("Working") 

133 time.sleep(0.1) 

134 spinner.update_value("Finishing") 

135 time.sleep(0.1) 

136 

137 output = string_io.getvalue() 

138 print(output) 

139 assert "Status: Working" in output 

140 assert "Status: Finishing" in output 

141 

142 

143def test_spinner_context_exception_handling(): 

144 string_io = io.StringIO() 

145 try: 

146 with SpinnerContext(output_stream=string_io, update_interval=0.05) as spinner: 

147 spinner.update_value("Before exception") 

148 time.sleep(0.1) 

149 raise ValueError("Test exception") 

150 except ValueError: 

151 pass 

152 

153 output = string_io.getvalue() 

154 print(output) 

155 assert "Before exception" in output 

156 assert output.endswith("\n") 

157 

158 

159def test_spinner_long_running_task(): 

160 string_io = io.StringIO() 

161 

162 @spinner_decorator( 

163 message="Iteration: ", 

164 output_stream=string_io, 

165 update_interval=0.05, 

166 mutable_kwarg_key="update", 

167 ) 

168 def long_task(iterations, update): 

169 for i in range(iterations): 

170 update(i + 1) 

171 time.sleep(0.1) 

172 

173 long_task(3, update=lambda x: x) 

174 

175 output = string_io.getvalue() 

176 print(output) 

177 for i in range(1, 4): 

178 assert f"Iteration: {i}" in output