Coverage for tests/test_augment.py: 100%

1import numpy as np 

2import pytest 

3from evutils.augment import drop_random_events 

4from evutils.types import Event_dtype 

5 

6def test_drop_random_events(): 

7 events = np.array( 

8 [(i, i, i*100, 1) for i in range(100)], 

9 dtype=Event_dtype 

10 ) 

11 

12 # Normal case 

13 dropped = drop_random_events(events, drop_rate=0.1) 

14 assert len(dropped) == 90 

15 

16 # Value errors 

17 with pytest.raises(ValueError, match="drop_rate must be between 0 and 1"): 

18 drop_random_events(events, drop_rate=0.0) 

19 

20 with pytest.raises(ValueError, match="drop_rate must be between 0 and 1"): 

21 drop_random_events(events, drop_rate=1.0) 

22 

23 with pytest.raises(ValueError, match="drop_rate must be between 0 and 1"): 

24 drop_random_events(events, drop_rate=-0.1) 

25 

26 with pytest.raises(ValueError, match="drop_rate must be between 0 and 1"): 

27 drop_random_events(events, drop_rate=1.5) 

28 

29 # Empty array 

30 empty_events = np.array([], dtype=Event_dtype) 

31 assert len(drop_random_events(empty_events, drop_rate=0.5)) == 0