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}
|