Coverage for intelligence_toolkit/tests/unit/helpers/test_progress_batch_callback.py: 100%

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

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

3# 

4import pytest 

5 

6from intelligence_toolkit.helpers.progress_batch_callback import ProgressBatchCallback 

7 

8 

9def test_progress_batch_callback_initialization(): 

10 callback = ProgressBatchCallback() 

11 

12 assert callback.current_batch == 0 

13 assert callback.total_batches == 0 

14 

15 

16def test_progress_batch_callback_on_batch_change(): 

17 callback = ProgressBatchCallback() 

18 

19 callback.on_batch_change(5, 10) 

20 

21 assert callback.current_batch == 5 

22 assert callback.total_batches == 10 

23 

24 

25def test_progress_batch_callback_on_batch_change_with_message(): 

26 callback = ProgressBatchCallback() 

27 

28 callback.on_batch_change(3, 7, message="Processing...") 

29 

30 assert callback.current_batch == 3 

31 assert callback.total_batches == 7 

32 assert callback.message == "Processing..." 

33 

34 

35def test_progress_batch_callback_multiple_updates(): 

36 callback = ProgressBatchCallback() 

37 

38 callback.on_batch_change(1, 10) 

39 assert callback.current_batch == 1 

40 

41 callback.on_batch_change(5, 10) 

42 assert callback.current_batch == 5 

43 

44 callback.on_batch_change(10, 10) 

45 assert callback.current_batch == 10 

46 

47 

48def test_progress_batch_callback_message_attribute_exists(): 

49 callback = ProgressBatchCallback() 

50 callback.on_batch_change(1, 1, message="test") 

51 

52 assert hasattr(callback, 'message') 

53 assert callback.message == "test" 

54 

55 

56def test_progress_batch_callback_default_message(): 

57 callback = ProgressBatchCallback() 

58 callback.on_batch_change(1, 1) 

59 

60 # Default message is empty string 

61 assert callback.message == ""