Coverage for src/tests/core/test_functions.py: 100%

1"""Module for testing the functions in functions.py.""" 

2 

3from kwai_core.functions import async_groupby, generate_filenames 

4 

5 

6async def test_async_groupby(): 

7 """Test async_groupby method.""" 

8 

9 async def _generate_group(): 

10 yield "A", 1 

11 yield "A", 2 

12 yield "B", 3 

13 yield "C", 4 

14 yield "C", 5 

15 

16 it = async_groupby(_generate_group(), lambda x: x[0]) 

17 groups = {key: group async for key, group in it} 

18 

19 assert groups["A"] == [("A", 1), ("A", 2)] 

20 assert groups["B"] == [("B", 3)] 

21 assert groups["C"] == [("C", 4), ("C", 5)] 

22 

23 

24def test_generate_filenames(): 

25 """Test generate_filenames.""" 

26 filename_generator = generate_filenames("test_", ".csv") 

27 

28 assert next(filename_generator) == "test_001.csv" 

29 assert next(filename_generator) == "test_002.csv" 

30 assert next(filename_generator) == "test_003.csv" 

31 

32 

33def test_generate_filenames_with_shorter_suffix(): 

34 """Test generate_filenames.""" 

35 filename_generator = generate_filenames("test_", ".csv", 2) 

36 

37 assert next(filename_generator) == "test_01.csv" 

38 assert next(filename_generator) == "test_02.csv" 

39 assert next(filename_generator) == "test_03.csv"