1#!/usr/bin/env python3
2
3# Copyright 2020 The 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
17# Library to extract scenario definitions from scenario_config.py.
18#
19# Contains functions to filter, analyze and dump scenario definitions.
20#
21# This library is used in loadtest_config.py to generate the "scenariosJSON"
22# field in the format accepted by the OSS benchmarks framework.
23# See https://github.com/grpc/test-infra/blob/master/config/samples/cxx_example_loadtest.yaml
24#
25# It can also be used to dump scenarios to files, to count scenarios by
26# language, and to export scenario languages in a format that can be used for
27# automation.
28#
29# Example usage:
30#
31# scenario_config.py --export_scenarios -l cxx -f cxx_scenario_ -r '.*' \
32#     --category=scalable
33#
34# scenario_config.py --count_scenarios
35#
36# scenario_config.py --count_scenarios --category=scalable
37#
38# For usage of the language config output, see loadtest_config.py.
39
40import argparse
41import collections
42import json
43import re
44import sys
45from typing import Any, Callable, Dict, Iterable, NamedTuple
46
47import scenario_config
48
49# Language parameters for load test config generation.
50
51LanguageConfig = NamedTuple('LanguageConfig', [('category', str),
52                                               ('language', str),
53                                               ('client_language', str),
54                                               ('server_language', str)])
55
56
57def category_string(categories: Iterable[str], category: str) -> str:
58    """Converts a list of categories into a single string for counting."""
59    if category != 'all':
60        return category if category in categories else ''
61
62    main_categories = ('scalable', 'smoketest')
63    s = set(categories)
64
65    c = [m for m in main_categories if m in s]
66    s.difference_update(main_categories)
67    c.extend(s)
68    return ' '.join(c)
69
70
71def gen_scenario_languages(category: str) -> Iterable[LanguageConfig]:
72    """Generates tuples containing the languages specified in each scenario."""
73    for language in scenario_config.LANGUAGES:
74        for scenario in scenario_config.LANGUAGES[language].scenarios():
75            client_language = scenario.get('CLIENT_LANGUAGE', '')
76            server_language = scenario.get('SERVER_LANGUAGE', '')
77            categories = scenario.get('CATEGORIES', [])
78            if category != 'all' and category not in categories:
79                continue
80            cat = category_string(categories, category)
81            yield LanguageConfig(category=cat,
82                                 language=language,
83                                 client_language=client_language,
84                                 server_language=server_language)
85
86
87def scenario_filter(
88    scenario_name_regex: str = '.*',
89    category: str = 'all',
90    client_language: str = '',
91    server_language: str = '',
92) -> Callable[[Dict[str, Any]], bool]:
93    """Returns a function to filter scenarios to process."""
94
95    def filter_scenario(scenario: Dict[str, Any]) -> bool:
96        """Filters scenarios that match specified criteria."""
97        if not re.search(scenario_name_regex, scenario["name"]):
98            return False
99        # if the 'CATEGORIES' key is missing, treat scenario as part of
100        # 'scalable' and 'smoketest'. This matches the behavior of
101        # run_performance_tests.py.
102        scenario_categories = scenario.get('CATEGORIES',
103                                           ['scalable', 'smoketest'])
104        if category not in scenario_categories and category != 'all':
105            return False
106
107        scenario_client_language = scenario.get('CLIENT_LANGUAGE', '')
108        if client_language != scenario_client_language:
109            return False
110
111        scenario_server_language = scenario.get('SERVER_LANGUAGE', '')
112        if server_language != scenario_server_language:
113            return False
114
115        return True
116
117    return filter_scenario
118
119
120def gen_scenarios(
121    language_name: str, scenario_filter_function: Callable[[Dict[str, Any]],
122                                                           bool]
123) -> Iterable[Dict[str, Any]]:
124    """Generates scenarios that match a given filter function."""
125    return map(
126        scenario_config.remove_nonproto_fields,
127        filter(scenario_filter_function,
128               scenario_config.LANGUAGES[language_name].scenarios()))
129
130
131def dump_to_json_files(scenarios: Iterable[Dict[str, Any]],
132                       filename_prefix: str) -> None:
133    """Dumps a list of scenarios to JSON files"""
134    count = 0
135    for scenario in scenarios:
136        filename = '{}{}.json'.format(filename_prefix, scenario['name'])
137        print('Writing file {}'.format(filename), file=sys.stderr)
138        with open(filename, 'w') as outfile:
139            # The dump file should have {"scenarios" : []} as the top level
140            # element, when embedded in a LoadTest configuration YAML file.
141            json.dump({'scenarios': [scenario]}, outfile, indent=2)
142        count += 1
143    print('Wrote {} scenarios'.format(count), file=sys.stderr)
144
145
146def main() -> None:
147    language_choices = sorted(scenario_config.LANGUAGES.keys())
148    argp = argparse.ArgumentParser(description='Exports scenarios to files.')
149    argp.add_argument('--export_scenarios',
150                      action='store_true',
151                      help='Export scenarios to JSON files.')
152    argp.add_argument('--count_scenarios',
153                      action='store_true',
154                      help='Count scenarios for all test languages.')
155    argp.add_argument('-l',
156                      '--language',
157                      choices=language_choices,
158                      help='Language to export.')
159    argp.add_argument('-f',
160                      '--filename_prefix',
161                      default='scenario_dump_',
162                      type=str,
163                      help='Prefix for exported JSON file names.')
164    argp.add_argument('-r',
165                      '--regex',
166                      default='.*',
167                      type=str,
168                      help='Regex to select scenarios to run.')
169    argp.add_argument(
170        '--category',
171        default='all',
172        choices=['all', 'inproc', 'scalable', 'smoketest', 'sweep'],
173        help='Select scenarios for a category of tests.')
174    argp.add_argument(
175        '--client_language',
176        default='',
177        choices=language_choices,
178        help='Select only scenarios with a specified client language.')
179    argp.add_argument(
180        '--server_language',
181        default='',
182        choices=language_choices,
183        help='Select only scenarios with a specified server language.')
184    args = argp.parse_args()
185
186    if args.export_scenarios and not args.language:
187        print('Dumping scenarios requires a specified language.',
188              file=sys.stderr)
189        argp.print_usage(file=sys.stderr)
190        return
191
192    if args.export_scenarios:
193        s_filter = scenario_filter(scenario_name_regex=args.regex,
194                                   category=args.category,
195                                   client_language=args.client_language,
196                                   server_language=args.server_language)
197        scenarios = gen_scenarios(args.language, s_filter)
198        dump_to_json_files(scenarios, args.filename_prefix)
199
200    if args.count_scenarios:
201        print('Scenario count for all languages (category: {}):'.format(
202            args.category))
203        print('{:>5}  {:16} {:8} {:8} {}'.format('Count', 'Language', 'Client',
204                                                 'Server', 'Categories'))
205        c = collections.Counter(gen_scenario_languages(args.category))
206        total = 0
207        for ((cat, l, cl, sl), count) in c.most_common():
208            print('{count:5}  {l:16} {cl:8} {sl:8} {cat}'.format(l=l,
209                                                                 cl=cl,
210                                                                 sl=sl,
211                                                                 count=count,
212                                                                 cat=cat))
213            total += count
214
215        print('\n{:>5}  total scenarios (category: {})'.format(
216            total, args.category))
217
218
219if __name__ == "__main__":
220    main()
221