Coverage for intelligence_toolkit/tests/unit/helpers/test_decorators.py: 98%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2# Licensed under the MIT license. See LICENSE file in the project. 

3# 

4import time 

5from unittest.mock import MagicMock 

6 

7import pytest 

8 

9from intelligence_toolkit.helpers.decorators import retry_with_backoff 

10 

11 

12def test_retry_with_backoff_success_on_first_try(): 

13 mock_func = MagicMock(return_value="success") 

14 decorated = retry_with_backoff()(mock_func) 

15 

16 result = decorated() 

17 

18 assert result == "success" 

19 assert mock_func.call_count == 1 

20 

21 

22def test_retry_with_backoff_success_after_retries(): 

23 mock_func = MagicMock(side_effect=[Exception("Error 1"), Exception("Error 2"), "success"]) 

24 decorated = retry_with_backoff(retries=5, backoff_in_seconds=0.01)(mock_func) 

25 

26 result = decorated() 

27 

28 assert result == "success" 

29 assert mock_func.call_count == 3 

30 

31 

32def test_retry_with_backoff_max_retries_exceeded(): 

33 mock_func = MagicMock(side_effect=Exception("Persistent error")) 

34 decorated = retry_with_backoff(retries=3, backoff_in_seconds=0.01)(mock_func) 

35 

36 with pytest.raises(Exception, match="Persistent error"): 

37 decorated() 

38 

39 assert mock_func.call_count == 4 # Initial call + 3 retries 

40 

41 

42def test_retry_with_backoff_preserves_function_name(): 

43 def test_function(): 

44 return "test" 

45 

46 decorated = retry_with_backoff()(test_function) 

47 

48 assert decorated.__name__ == "test_function" 

49 

50 

51def test_retry_with_backoff_with_arguments(): 

52 mock_func = MagicMock(side_effect=[Exception("Error"), "success"]) 

53 decorated = retry_with_backoff(retries=3, backoff_in_seconds=0.01)(mock_func) 

54 

55 result = decorated("arg1", "arg2", kwarg1="value1") 

56 

57 assert result == "success" 

58 assert mock_func.call_count == 2 

59 mock_func.assert_called_with("arg1", "arg2", kwarg1="value1") 

60 

61 

62def test_retry_with_backoff_exponential_backoff(): 

63 mock_func = MagicMock(side_effect=[Exception("Error 1"), Exception("Error 2"), "success"]) 

64 

65 start_time = time.time() 

66 decorated = retry_with_backoff(retries=5, backoff_in_seconds=0.1)(mock_func) 

67 result = decorated() 

68 elapsed = time.time() - start_time 

69 

70 assert result == "success" 

71 # Should have some delay due to backoff (0.1 * 2^0 + 0.1 * 2^1 ≈ 0.3 seconds minimum) 

72 assert elapsed >= 0.1 

73 

74 

75def test_retry_with_backoff_no_retries(): 

76 mock_func = MagicMock(side_effect=Exception("Immediate failure")) 

77 decorated = retry_with_backoff(retries=0, backoff_in_seconds=0.01)(mock_func) 

78 

79 with pytest.raises(Exception, match="Immediate failure"): 

80 decorated() 

81 

82 assert mock_func.call_count == 1 

83 

84 

85def test_retry_with_backoff_returns_correct_value(): 

86 expected_value = {"key": "value", "number": 42} 

87 mock_func = MagicMock(return_value=expected_value) 

88 decorated = retry_with_backoff()(mock_func) 

89 

90 result = decorated() 

91 

92 assert result == expected_value