1#!/usr/bin/env python3
2
3# Copyright 2022 gRPC authors.
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17import os
18import re
19import subprocess
20import sys
21
22from google.cloud import bigquery
23import run_buildozer
24import update_flakes_query
25
26lookback_hours = 24 * 7 * 4
27
28
29def include_test(test):
30    if '@' in test:
31        return False
32    if test.startswith("//test/cpp/qps:"):
33        return False
34    return True
35
36
37TEST_DIRS = ['test/core', 'test/cpp']
38tests = {}
39already_flaky = set()
40for test_dir in TEST_DIRS:
41    for line in subprocess.check_output(
42        ['bazel', 'query', 'tests({}/...)'.format(test_dir)]).splitlines():
43        test = line.strip().decode('utf-8')
44        if not include_test(test):
45            continue
46        tests[test] = False
47for test_dir in TEST_DIRS:
48    for line in subprocess.check_output(
49        ['bazel', 'query',
50         'attr(flaky, 1, tests({}/...))'.format(test_dir)]).splitlines():
51        test = line.strip().decode('utf-8')
52        if not include_test(test):
53            continue
54        already_flaky.add(test)
55
56flaky_e2e = set()
57
58client = bigquery.Client()
59for row in client.query(
60        update_flakes_query.QUERY.format(
61            lookback_hours=lookback_hours)).result():
62    if "/macos/" in row.job_name:
63        continue  # we know mac stuff is flaky
64    if row.test_binary not in tests:
65        m = re.match(r'^//test/core/end2end:([^@]*)@([^@]*)(.*)',
66                     row.test_binary)
67        if m:
68            flaky_e2e.add('{}@{}{}'.format(m.group(1), m.group(2), m.group(3)))
69            print("will mark end2end test {} as flaky".format(row.test_binary))
70        else:
71            print("skip obsolete test {}".format(row.test_binary))
72        continue
73    print("will mark {} as flaky".format(row.test_binary))
74    tests[row.test_binary] = True
75
76buildozer_commands = []
77for test, flaky in sorted(tests.items()):
78    if flaky:
79        buildozer_commands.append('set flaky True|{}'.format(test))
80    elif test in already_flaky:
81        buildozer_commands.append('remove flaky|{}'.format(test))
82
83with open('test/core/end2end/flaky.bzl', 'w') as f:
84    with open(sys.argv[0]) as my_source:
85        for line in my_source:
86            if line[0] != '#':
87                break
88        for line in my_source:
89            if line[0] == '#':
90                print(line.strip(), file=f)
91                break
92        for line in my_source:
93            if line[0] != '#':
94                break
95            print(line.strip(), file=f)
96    print(
97        "\"\"\"A list of flaky tests, consumed by generate_tests.bzl to set flaky attrs.\"\"\"",
98        file=f)
99    print("FLAKY_TESTS = [", file=f)
100    for line in sorted(list(flaky_e2e)):
101        print("    \"{}\",".format(line), file=f)
102    print("]", file=f)
103
104run_buildozer.run_buildozer(buildozer_commands)
105