add support for scenario categories

hdszju · May 6, 2016 · a21c7e9 · a21c7e9
1 parent 541d5d7
commit a21c7e9
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 7 deletions.
diff --git a/test/cpp/qps/gen_build_yaml.py b/test/cpp/qps/gen_build_yaml.py
@@ -43,12 +43,16 @@
 
 import performance.scenario_config as scenario_config
 
+def _scenario_json_string(scenario_json):
+  return json.dumps(scenario_config.remove_nonproto_fields(scenario_json))
+
 print yaml.dump({
   'tests': [
     {
       'name': 'json_run_localhost',
-      'shortname': 'json_run_localhost:%s' % js['name'],
-      'args': ['--scenario_json', pipes.quote(json.dumps(js))],
+      'shortname': 'json_run_localhost:%s' % scenario_json['name'],
+      'args': ['--scenario_json',
+               pipes.quote(_scenario_json_string(scenario_json))],
       'ci_platforms': ['linux', 'mac', 'posix', 'windows'],
       'platforms': ['linux', 'mac', 'posix', 'windows'],
       'flaky': False,
@@ -58,6 +62,6 @@
       'cpu_cost': 1000.0,
       'exclude_configs': []
     }
-    for js in scenario_config.CXXLanguage().scenarios()
+    for scenario_json in scenario_config.CXXLanguage().scenarios()
   ]
 })
diff --git a/tools/run_tests/performance/scenario_config.py b/tools/run_tests/performance/scenario_config.py
@@ -76,6 +76,14 @@ def _get_secargs(is_secure):
     return None
 
 
+def remove_nonproto_fields(scenario):
+  """Remove special-purpose that contains some extra info about the scenario
+  but don't belong to the ScenarioConfig protobuf message"""
+  scenario.pop('CATEGORIES', None)
+  scenario.pop('SERVER_LANGUAGE', None)
+  return scenario
+
+
 def _ping_pong_scenario(name, rpc_type,
                         client_type, server_type,
                         secure=True,
@@ -84,7 +92,8 @@ def _ping_pong_scenario(name, rpc_type,
                         server_language=None,
                         server_core_limit=0,
                         async_server_threads=0,
-                        warmup_seconds=WARMUP_SECONDS):
+                        warmup_seconds=WARMUP_SECONDS,
+                        categories=[]):
   """Creates a basic ping pong scenario."""
   scenario = {
     'name': name,
@@ -135,6 +144,8 @@ def _ping_pong_scenario(name, rpc_type,
   if server_language:
     # the SERVER_LANGUAGE field is recognized by run_performance_tests.py
     scenario['SERVER_LANGUAGE'] = server_language
+  if categories:
+    scenario['CATEGORIES'] = categories
   return scenario
 
 

diff --git a/tools/run_tests/run_performance_tests.py b/tools/run_tests/run_performance_tests.py
@@ -255,9 +255,9 @@ def create_scenarios(languages, workers_by_lang, remote_host=None, regex='.*',
       if re.search(args.regex, scenario_json['name']):
         workers = workers_by_lang[str(language)]
         # 'SERVER_LANGUAGE' is an indicator for this script to pick
-        # a server in different language. It doesn't belong to the Scenario
-        # schema, so we also need to remove it.
-        custom_server_lang = scenario_json.pop('SERVER_LANGUAGE', None)
+        # a server in different language.
+        custom_server_lang = scenario_json.get('SERVER_LANGUAGE', None)
+        scenario_json = scenario_config.remove_nonproto_fields(scenario_json)
         if custom_server_lang:
           if not workers_by_lang.get(custom_server_lang, []):
             print 'Warning: Skipping scenario %s as' % scenario_json['name']