aboutsummaryrefslogtreecommitdiffstats
path: root/library/python/testing/yatest_lib/test_splitter.py
blob: ba90a47291c4c97658e3823c5a0749988d7cea84 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
# coding: utf-8
 
import collections 
 
 
def flatten_tests(test_classes):
    """ 
    >>> test_classes = {x: [x] for x in range(5)} 
    >>> flatten_tests(test_classes)
    [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)] 
    >>> test_classes = {x: [x + 1, x + 2] for x in range(2)} 
    >>> flatten_tests(test_classes)
    [(0, 1), (0, 2), (1, 2), (1, 3)] 
    """ 
    tests = [] 
    for class_name, test_names in test_classes.items(): 
        tests += [(class_name, test_name) for test_name in test_names] 
    return tests 
 
 
def get_sequential_chunk(tests, modulo, modulo_index, is_sorted=False):
    """ 
    >>> get_sequential_chunk(range(10), 4, 0)
    [0, 1, 2] 
    >>> get_sequential_chunk(range(10), 4, 1)
    [3, 4, 5] 
    >>> get_sequential_chunk(range(10), 4, 2)
    [6, 7]
    >>> get_sequential_chunk(range(10), 4, 3)
    [8, 9]
    >>> get_sequential_chunk(range(10), 4, 4)
    [] 
    >>> get_sequential_chunk(range(10), 4, 5)
    [] 
    """ 
    if not is_sorted:
        tests = sorted(tests)
    chunk_size = len(tests) // modulo
    not_used = len(tests) % modulo
    shift = chunk_size + (modulo_index < not_used)
    start = chunk_size * modulo_index + min(modulo_index, not_used)
    end = start + shift 
    return [] if end > len(tests) else tests[start:end]
 
 
def get_shuffled_chunk(tests, modulo, modulo_index, is_sorted=False):
    """ 
    >>> get_shuffled_chunk(range(10), 4, 0)
    [0, 4, 8]
    >>> get_shuffled_chunk(range(10), 4, 1)
    [1, 5, 9]
    >>> get_shuffled_chunk(range(10), 4, 2)
    [2, 6]
    >>> get_shuffled_chunk(range(10), 4, 3)
    [3, 7]
    >>> get_shuffled_chunk(range(10), 4, 4)
    []
    >>> get_shuffled_chunk(range(10), 4, 5)
    []
    """
    if not is_sorted:
        tests = sorted(tests)
    result_tests = []
    for i, test in enumerate(tests):
        if i % modulo == modulo_index:
            result_tests.append(test)
    return result_tests


def get_splitted_tests(test_entities, modulo, modulo_index, partition_mode, is_sorted=False):
    if partition_mode == 'SEQUENTIAL':
        return get_sequential_chunk(test_entities, modulo, modulo_index, is_sorted)
    elif partition_mode == 'MODULO':
        return get_shuffled_chunk(test_entities, modulo, modulo_index, is_sorted)
    else:
        raise ValueError("detected unknown partition mode: {}".format(partition_mode))


def filter_tests_by_modulo(test_classes, modulo, modulo_index, split_by_tests, partition_mode="SEQUENTIAL"):
    """
    >>> test_classes = {x: [x] for x in range(20)} 
    >>> filter_tests_by_modulo(test_classes, 4, 0, False) 
    {0: [0], 1: [1], 2: [2], 3: [3], 4: [4]} 
    >>> filter_tests_by_modulo(test_classes, 4, 1, False) 
    {8: [8], 9: [9], 5: [5], 6: [6], 7: [7]} 
    >>> filter_tests_by_modulo(test_classes, 4, 2, False) 
    {10: [10], 11: [11], 12: [12], 13: [13], 14: [14]} 
 
    >>> dict(filter_tests_by_modulo(test_classes, 4, 0, True)) 
    {0: [0], 1: [1], 2: [2], 3: [3], 4: [4]} 
    >>> dict(filter_tests_by_modulo(test_classes, 4, 1, True)) 
    {8: [8], 9: [9], 5: [5], 6: [6], 7: [7]} 
    """ 
    if split_by_tests: 
        tests = get_splitted_tests(flatten_tests(test_classes), modulo, modulo_index, partition_mode)
        test_classes = collections.defaultdict(list) 
        for class_name, test_name in tests: 
            test_classes[class_name].append(test_name) 
        return test_classes 
    else: 
        target_classes = get_splitted_tests(test_classes, modulo, modulo_index, partition_mode)
        return {class_name: test_classes[class_name] for class_name in target_classes}