1#!/usr/bin/env vpython3 2# Copyright 2018 The Chromium Authors 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5"""Tests for perf_device_trigger_unittest.py.""" 6 7import unittest 8 9import perf_device_trigger 10 11 12class Args(object): # pylint: disable=useless-object-inheritance 13 def __init__(self): 14 self.shards = 1 15 self.shard_index = None 16 self.dump_json = '' 17 self.multiple_trigger_configs = None 18 self.multiple_dimension_script_verbose = False 19 self.use_dynamic_shards = False 20 21 22class FakeTriggerer(perf_device_trigger.PerfDeviceTriggerer): 23 def __init__(self, args, swarming_args, files, list_bots_result, 24 list_tasks_results): 25 self._bot_statuses = [] 26 self._swarming_runs = [] 27 self._files = files 28 self._temp_file_id = 0 29 self._triggered_with_swarming_go = 0 30 self._list_bots_result = list_bots_result 31 self._list_tasks_results = list_tasks_results 32 # pylint: disable=super-with-arguments 33 super(FakeTriggerer, self).__init__(args, swarming_args) 34 # pylint: enable=super-with-arguments 35 36 def set_files(self, files): 37 self._files = files 38 39 def make_temp_file(self, prefix=None, suffix=None): 40 result = prefix + str(self._temp_file_id) + suffix 41 self._temp_file_id += 1 42 return result 43 44 def delete_temp_file(self, temp_file): 45 pass 46 47 def read_json_from_temp_file(self, temp_file): 48 return self._files[temp_file] 49 50 def read_encoded_json_from_temp_file(self, temp_file): 51 return self._files[temp_file] 52 53 def write_json_to_file(self, merged_json, output_file): 54 self._files[output_file] = merged_json 55 56 def list_bots(self, 57 dimensions, 58 server='chromium-swarm.appspot.com'): 59 return self._list_bots_result 60 61 def list_tasks(self, tags, limit=None, 62 server='chromium-swarm.appspot.com'): 63 res, self._list_tasks_results = self._list_tasks_results[ 64 0], self._list_tasks_results[1:] 65 return res 66 67 def run_swarming(self, args): 68 self._swarming_runs.append(args) 69 70 def run_swarming_go(self, 71 args, 72 _json_path, 73 _shard_index, 74 _shard, 75 _merged_json=None): 76 self._triggered_with_swarming_go += 1 77 self.run_swarming(args) 78 79 80class UnitTest(unittest.TestCase): 81 def setup_and_trigger(self, 82 previous_task_assignment_map, 83 alive_bots, 84 dead_bots, 85 use_dynamic_shards=False): 86 args = Args() 87 args.shards = len(previous_task_assignment_map) 88 args.dump_json = 'output.json' 89 args.multiple_dimension_script_verbose = True 90 if use_dynamic_shards: 91 args.use_dynamic_shards = True 92 swarming_args = [ 93 'trigger', 94 '--swarming', 95 'http://foo_server', 96 '--dimension', 97 'pool', 98 'chrome-perf-fyi', 99 '--dimension', 100 'os', 101 'windows', 102 '--', 103 'benchmark1', 104 ] 105 106 triggerer = FakeTriggerer( 107 args, swarming_args, self.get_files(args.shards), 108 self.generate_list_of_eligible_bots_query_response( 109 alive_bots, dead_bots), [ 110 self.generate_last_task_to_shard_query_response( 111 i, previous_task_assignment_map.get(i)) 112 for i in range(args.shards) 113 ]) 114 triggerer.trigger_tasks(args, swarming_args) 115 return triggerer 116 117 def get_files(self, num_shards): 118 files = {} 119 file_index = 0 120 file_index = file_index + 1 121 # Perf device trigger will call swarming n times: 122 # 1. Once for all eligible bots 123 # 2. once per shard to determine last bot run 124 # Shard builders is a list of build ids that represents 125 # the last build that ran the shard that corresponds to that 126 # index. If that shard hasn't been run before the entry 127 # should be an empty string. 128 for i in range(num_shards): 129 task = { 130 'tasks': [{ 131 'request': { 132 'task_id': 'f%d' % i, 133 }, 134 }], 135 } 136 files['base_trigger_dimensions%d.json' % file_index] = task 137 file_index = file_index + 1 138 return files 139 140 def generate_last_task_to_shard_query_response(self, shard, bot_id): 141 if len(bot_id): 142 # Test both cases where bot_id is present and you have to parse 143 # out of the tags. 144 if shard % 2: 145 return [{'bot_id': bot_id}] 146 return [{'tags': ['id:%s' % bot_id]}] 147 return [] 148 149 def generate_list_of_eligible_bots_query_response(self, alive_bots, 150 dead_bots): 151 if len(alive_bots) == 0 and len(dead_bots) == 0: 152 return {} 153 bots = [] 154 for bot_id in alive_bots: 155 bots.append({ 156 'bot_id': ('%s' % bot_id), 157 'is_dead': False, 158 'quarantined': False 159 }) 160 is_dead = True 161 for bot_id in dead_bots: 162 is_quarantined = (not is_dead) 163 bots.append({ 164 'bot_id': ('%s' % bot_id), 165 'is_dead': is_dead, 166 'quarantined': is_quarantined 167 }) 168 is_dead = (not is_dead) 169 return bots 170 171 def list_contains_sublist(self, main_list, sub_list): 172 return any(sub_list == main_list[offset:offset + len(sub_list)] 173 for offset in range(len(main_list) - (len(sub_list) - 1))) 174 175 def get_triggered_shard_to_bot(self, triggerer): 176 triggered_map = {} 177 for run in triggerer._swarming_runs: 178 if not 'trigger' in run: 179 continue 180 bot_id = run[(run.index('id') + 1)] 181 182 g = 'GTEST_SHARD_INDEX=' 183 shard = [int(r[len(g):]) for r in run if r.startswith(g)][0] 184 185 triggered_map[shard] = bot_id 186 return triggered_map 187 188 def test_all_healthy_shards(self): 189 triggerer = self.setup_and_trigger( 190 previous_task_assignment_map={ 191 0: 'build3', 192 1: 'build4', 193 2: 'build5' 194 }, 195 alive_bots=['build3', 'build4', 'build5'], 196 dead_bots=['build1', 'build2']) 197 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 198 self.assertEquals(len(set(expected_task_assignment.values())), 3) 199 200 # All three bots were healthy so we should expect the task assignment to 201 # stay the same 202 self.assertEquals(expected_task_assignment.get(0), 'build3') 203 self.assertEquals(expected_task_assignment.get(1), 'build4') 204 self.assertEquals(expected_task_assignment.get(2), 'build5') 205 206 def test_no_bot_returned(self): 207 with self.assertRaises(ValueError) as context: 208 self.setup_and_trigger(previous_task_assignment_map={0: 'build1'}, 209 alive_bots=[], 210 dead_bots=[]) 211 err_msg = 'Not enough available machines exist in swarming pool' 212 self.assertTrue(err_msg in str(context.exception)) 213 214 def test_previously_healthy_now_dead(self): 215 # Test that it swaps out build1 and build2 that are dead 216 # for two healthy bots 217 triggerer = self.setup_and_trigger( 218 previous_task_assignment_map={ 219 0: 'build1', 220 1: 'build2', 221 2: 'build3' 222 }, 223 alive_bots=['build3', 'build4', 'build5'], 224 dead_bots=['build1', 'build2']) 225 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 226 self.assertEquals(len(set(expected_task_assignment.values())), 3) 227 228 # The first two should be assigned to one of the unassigned healthy bots 229 new_healthy_bots = ['build4', 'build5'] 230 self.assertIn(expected_task_assignment.get(0), new_healthy_bots) 231 self.assertIn(expected_task_assignment.get(1), new_healthy_bots) 232 self.assertEquals(expected_task_assignment.get(2), 'build3') 233 234 def test_not_enough_healthy_bots(self): 235 triggerer = self.setup_and_trigger( 236 previous_task_assignment_map={ 237 0: 'build1', 238 1: 'build2', 239 2: 'build3', 240 3: 'build4', 241 4: 'build5' 242 }, 243 alive_bots=['build3', 'build4', 'build5'], 244 dead_bots=['build1', 'build2']) 245 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 246 self.assertEquals(len(set(expected_task_assignment.values())), 5) 247 248 # We have 5 shards and 5 bots that ran them, but two 249 # are now dead and there aren't any other healthy bots 250 # to swap out to. Make sure they still assign to the 251 # same shards. 252 self.assertEquals(expected_task_assignment.get(0), 'build1') 253 self.assertEquals(expected_task_assignment.get(1), 'build2') 254 self.assertEquals(expected_task_assignment.get(2), 'build3') 255 self.assertEquals(expected_task_assignment.get(3), 'build4') 256 self.assertEquals(expected_task_assignment.get(4), 'build5') 257 258 def test_not_enough_healthy_bots_shard_not_seen(self): 259 triggerer = self.setup_and_trigger( 260 previous_task_assignment_map={ 261 0: 'build1', 262 1: '', 263 2: 'build3', 264 3: 'build4', 265 4: 'build5' 266 }, 267 alive_bots=['build3', 'build4', 'build5'], 268 dead_bots=['build1', 'build2']) 269 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 270 self.assertEquals(len(set(expected_task_assignment.values())), 5) 271 272 # Not enough healthy bots so make sure shard 0 is still assigned to its 273 # same dead bot. 274 self.assertEquals(expected_task_assignment.get(0), 'build1') 275 # Shard 1 had not been triggered yet, but there weren't enough 276 # healthy bots. Make sure it got assigned to the other dead bot. 277 self.assertEquals(expected_task_assignment.get(1), 'build2') 278 # The rest of the assignments should stay the same. 279 self.assertEquals(expected_task_assignment.get(2), 'build3') 280 self.assertEquals(expected_task_assignment.get(3), 'build4') 281 self.assertEquals(expected_task_assignment.get(4), 'build5') 282 283 def test_shards_not_triggered_yet(self): 284 # First time this configuration has been seen. Choose three 285 # healthy shards to trigger jobs on 286 triggerer = self.setup_and_trigger( 287 previous_task_assignment_map={ 288 0: '', 289 1: '', 290 2: '' 291 }, 292 alive_bots=['build3', 'build4', 'build5'], 293 dead_bots=['build1', 'build2']) 294 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 295 self.assertEquals(len(set(expected_task_assignment.values())), 3) 296 new_healthy_bots = ['build3', 'build4', 'build5'] 297 self.assertIn(expected_task_assignment.get(0), new_healthy_bots) 298 self.assertIn(expected_task_assignment.get(1), new_healthy_bots) 299 self.assertIn(expected_task_assignment.get(2), new_healthy_bots) 300 301 def test_previously_duplicate_task_assignments(self): 302 triggerer = self.setup_and_trigger( 303 previous_task_assignment_map={ 304 0: 'build3', 305 1: 'build3', 306 2: 'build5', 307 3: 'build6' 308 }, 309 alive_bots=['build3', 'build4', 'build5', 'build7'], 310 dead_bots=['build1', 'build6']) 311 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 312 313 # Test that the new assignment will add a new bot to avoid 314 # assign 'build3' to both shard 0 & shard 1 as before. 315 # It also replaces the dead 'build6' bot. 316 self.assertEquals(set(expected_task_assignment.values()), 317 {'build3', 'build4', 'build5', 'build7'}) 318 319 def test_dynamic_sharding(self): 320 triggerer = self.setup_and_trigger( 321 # The previous map should not matter. 322 previous_task_assignment_map={ 323 0: 'build301', 324 1: 'build1--', 325 2: 'build-blah' 326 }, 327 alive_bots=['build1', 'build2', 'build3', 'build4', 'build5'], 328 dead_bots=[], 329 use_dynamic_shards=True) 330 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 331 332 self.assertEquals(set(expected_task_assignment.values()), 333 {'build1', 'build2', 'build3', 'build4', 'build5'}) 334 335 def test_dynamic_sharding_with_dead_bots(self): 336 triggerer = self.setup_and_trigger( 337 # The previous map should not matter. 338 previous_task_assignment_map={ 339 0: 'build301', 340 1: 'build1--', 341 2: 'build-blah' 342 }, 343 alive_bots=['build2', 'build5', 'build3'], 344 dead_bots=['build1', 'build4'], 345 use_dynamic_shards=True) 346 expected_task_assignment = self.get_triggered_shard_to_bot(triggerer) 347 348 self.assertEquals(set(expected_task_assignment.values()), 349 {'build2', 'build3', 'build5'}) 350 351 352if __name__ == '__main__': 353 unittest.main() 354