1#!/usr/bin/env python
2
3from __future__ import absolute_import
4
5import os
6import random
7from collections import defaultdict
8from itertools import chain
9from unittest import TestCase
10
11from six.moves import range
12from six import iteritems
13
14import mozunit
15from manifestparser.filters import (
16    chunk_by_dir,
17    chunk_by_runtime,
18    chunk_by_slice,
19)
20
21here = os.path.dirname(os.path.abspath(__file__))
22
23
24class ChunkBySlice(TestCase):
25    """Test chunking related filters"""
26
27    def generate_tests(self, num, disabled=None):
28        disabled = disabled or []
29        tests = []
30        for i in range(num):
31            test = {"name": "test%i" % i}
32            if i in disabled:
33                test["disabled"] = ""
34            tests.append(test)
35        return tests
36
37    def run_all_combos(self, num_tests, disabled=None):
38        tests = self.generate_tests(num_tests, disabled=disabled)
39
40        for total in range(1, num_tests + 1):
41            res = []
42            res_disabled = []
43            for chunk in range(1, total + 1):
44                f = chunk_by_slice(chunk, total)
45                res.append(list(f(tests, {})))
46                if disabled:
47                    f.disabled = True
48                    res_disabled.append(list(f(tests, {})))
49
50            lengths = [len([t for t in c if "disabled" not in t]) for c in res]
51            # the chunk with the most tests should have at most one more test
52            # than the chunk with the least tests
53            self.assertLessEqual(max(lengths) - min(lengths), 1)
54
55            # chaining all chunks back together should equal the original list
56            # of tests
57            self.assertEqual(list(chain.from_iterable(res)), list(tests))
58
59            if disabled:
60                lengths = [len(c) for c in res_disabled]
61                self.assertLessEqual(max(lengths) - min(lengths), 1)
62                self.assertEqual(list(chain.from_iterable(res_disabled)), list(tests))
63
64    def test_chunk_by_slice(self):
65        chunk = chunk_by_slice(1, 1)
66        self.assertEqual(list(chunk([], {})), [])
67
68        self.run_all_combos(num_tests=1)
69        self.run_all_combos(num_tests=10, disabled=[1, 2])
70
71        num_tests = 67
72        disabled = list(i for i in range(num_tests) if i % 4 == 0)
73        self.run_all_combos(num_tests=num_tests, disabled=disabled)
74
75    def test_two_times_more_chunks_than_tests(self):
76        # test case for bug 1182817
77        tests = self.generate_tests(5)
78
79        total_chunks = 10
80        for i in range(1, total_chunks + 1):
81            # ensure IndexError is not raised
82            chunk_by_slice(i, total_chunks)(tests, {})
83
84
85class ChunkByDir(TestCase):
86    """Test chunking related filters"""
87
88    def generate_tests(self, dirs):
89        """
90        :param dirs: dict of the form,
91                        { <dir>: <num tests> }
92        """
93        i = 0
94        for d, num in iteritems(dirs):
95            for _ in range(num):
96                i += 1
97                name = "test%i" % i
98                test = {"name": name, "relpath": os.path.join(d, name)}
99                yield test
100
101    def run_all_combos(self, dirs):
102        tests = list(self.generate_tests(dirs))
103
104        deepest = max(len(t["relpath"].split(os.sep)) - 1 for t in tests)
105        for depth in range(1, deepest + 1):
106
107            def num_groups(tests):
108                unique = set()
109                for p in [t["relpath"] for t in tests]:
110                    p = p.split(os.sep)
111                    p = p[: min(depth, len(p) - 1)]
112                    unique.add(os.sep.join(p))
113                return len(unique)
114
115            for total in range(1, num_groups(tests) + 1):
116                res = []
117                for this in range(1, total + 1):
118                    f = chunk_by_dir(this, total, depth)
119                    res.append(list(f(tests, {})))
120
121                lengths = list(map(num_groups, res))
122                # the chunk with the most dirs should have at most one more
123                # dir than the chunk with the least dirs
124                self.assertLessEqual(max(lengths) - min(lengths), 1)
125
126                all_chunks = list(chain.from_iterable(res))
127                # chunk_by_dir will mess up order, but chained chunks should
128                # contain all of the original tests and be the same length
129                self.assertEqual(len(all_chunks), len(tests))
130                for t in tests:
131                    self.assertIn(t, all_chunks)
132
133    def test_chunk_by_dir(self):
134        chunk = chunk_by_dir(1, 1, 1)
135        self.assertEqual(list(chunk([], {})), [])
136
137        dirs = {
138            "a": 2,
139        }
140        self.run_all_combos(dirs)
141
142        dirs = {
143            "": 1,
144            "foo": 1,
145            "bar": 0,
146            "/foobar": 1,
147        }
148        self.run_all_combos(dirs)
149
150        dirs = {
151            "a": 1,
152            "b": 1,
153            "a/b": 2,
154            "a/c": 1,
155        }
156        self.run_all_combos(dirs)
157
158        dirs = {
159            "a": 5,
160            "a/b": 4,
161            "a/b/c": 7,
162            "a/b/c/d": 1,
163            "a/b/c/e": 3,
164            "b/c": 2,
165            "b/d": 5,
166            "b/d/e": 6,
167            "c": 8,
168            "c/d/e/f/g/h/i/j/k/l": 5,
169            "c/d/e/f/g/i/j/k/l/m/n": 2,
170            "c/e": 1,
171        }
172        self.run_all_combos(dirs)
173
174
175class ChunkByRuntime(TestCase):
176    """Test chunking related filters"""
177
178    def generate_tests(self, dirs):
179        """
180        :param dirs: dict of the form,
181                     { <dir>: <num tests> }
182        """
183        i = 0
184        for d, num in iteritems(dirs):
185            for _ in range(num):
186                i += 1
187                name = "test%i" % i
188                manifest = os.path.join(d, "manifest.ini")
189                test = {
190                    "name": name,
191                    "relpath": os.path.join(d, name),
192                    "manifest": manifest,
193                    "manifest_relpath": manifest,
194                }
195                yield test
196
197    def get_runtimes(self, tests):
198        runtimes = defaultdict(int)
199        for test in tests:
200            runtimes[test["manifest_relpath"]] += random.randint(0, 100)
201        return runtimes
202
203    def chunk_by_round_robin(self, tests, total, runtimes):
204        tests_by_manifest = []
205        for manifest, runtime in iteritems(runtimes):
206            mtests = [t for t in tests if t["manifest_relpath"] == manifest]
207            tests_by_manifest.append((runtime, mtests))
208        tests_by_manifest.sort(key=lambda x: x[0], reverse=False)
209
210        chunks = [[] for i in range(total)]
211        d = 1  # direction
212        i = 0
213        for runtime, batch in tests_by_manifest:
214            chunks[i].extend(batch)
215
216            # "draft" style (last pick goes first in the next round)
217            if (i == 0 and d == -1) or (i == total - 1 and d == 1):
218                d = -d
219            else:
220                i += d
221
222        # make sure this test algorithm is valid
223        all_chunks = list(chain.from_iterable(chunks))
224        self.assertEqual(len(all_chunks), len(tests))
225        for t in tests:
226            self.assertIn(t, all_chunks)
227        return chunks
228
229    def run_all_combos(self, dirs):
230        tests = list(self.generate_tests(dirs))
231        runtimes = self.get_runtimes(tests)
232
233        for total in range(1, len(dirs) + 1):
234            chunks = []
235            for this in range(1, total + 1):
236                f = chunk_by_runtime(this, total, runtimes)
237                ret = list(f(tests, {}))
238                chunks.append(ret)
239
240            # chunk_by_runtime will mess up order, but chained chunks should
241            # contain all of the original tests and be the same length
242            all_chunks = list(chain.from_iterable(chunks))
243            self.assertEqual(len(all_chunks), len(tests))
244            for t in tests:
245                self.assertIn(t, all_chunks)
246
247            # calculate delta between slowest and fastest chunks
248            def runtime_delta(chunks):
249                totals = []
250                for chunk in chunks:
251                    manifests = set([t["manifest_relpath"] for t in chunk])
252                    total = sum(runtimes[m] for m in manifests)
253                    totals.append(total)
254                return max(totals) - min(totals)
255
256            delta = runtime_delta(chunks)
257
258            # redo the chunking a second time using a round robin style
259            # algorithm
260            chunks = self.chunk_by_round_robin(tests, total, runtimes)
261            # sanity check the round robin algorithm
262            all_chunks = list(chain.from_iterable(chunks))
263            self.assertEqual(len(all_chunks), len(tests))
264            for t in tests:
265                self.assertIn(t, all_chunks)
266
267            # since chunks will never have exactly equal runtimes, it's hard
268            # to tell if they were chunked optimally. Make sure it at least
269            # beats a naive round robin approach.
270            self.assertLessEqual(delta, runtime_delta(chunks))
271
272    def test_chunk_by_runtime(self):
273        random.seed(42)
274
275        chunk = chunk_by_runtime(1, 1, {})
276        self.assertEqual(list(chunk([], {})), [])
277
278        dirs = {
279            "a": 2,
280        }
281        self.run_all_combos(dirs)
282
283        dirs = {
284            "": 1,
285            "foo": 1,
286            "bar": 0,
287            "/foobar": 1,
288        }
289        self.run_all_combos(dirs)
290
291        dirs = {
292            "a": 1,
293            "b": 1,
294            "a/b": 2,
295            "a/c": 1,
296        }
297        self.run_all_combos(dirs)
298
299        dirs = {
300            "a": 5,
301            "a/b": 4,
302            "a/b/c": 7,
303            "a/b/c/d": 1,
304            "a/b/c/e": 3,
305            "b/c": 2,
306            "b/d": 5,
307            "b/d/e": 6,
308            "c": 8,
309            "c/d/e/f/g/h/i/j/k/l": 5,
310            "c/d/e/f/g/i/j/k/l/m/n": 2,
311            "c/e": 1,
312        }
313        self.run_all_combos(dirs)
314
315
316if __name__ == "__main__":
317    mozunit.main()
318