1# Copyright 2015 The Chromium OS Authors. All rights reserved. 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4 5import collections 6import logging 7import multiprocessing 8 9from autotest_lib.client.common_lib import error 10from autotest_lib.client.common_lib.cros.network import xmlrpc_datatypes 11from autotest_lib.server.cros.network import connection_worker 12 13"""DUT Control module is used to control all the DUT's in a Clique set. 14We need to execute a sequence of steps on each DUT in the pool parallely and 15collect the results from all the executions. 16 17Class Hierarchy: 18---------------- 19 CliqueDUTControl 20 | 21 ------------------------------------------------------- 22 | | 23 CliqueDUTRole CliqueDUTBatch 24 | | 25 ------------------------------------- --------------------- 26 | | | | 27 DUTRoleConnectDisconnect DUTRoleFileTransfer CliqueDUTSet CliqueDUTPool 28 29CliqueDUTControl - Base control class. Stores and retrieves test params used 30for all control operations. Should never be directly instantiated. 31 32CliqueDUTRole - Used to control one single DUT in the test. This is a base class 33which should be derived to define a role to be performed by the DUT. Should 34never be directly instantiated. 35 36CliqueDUTBatch - Used to control a batch of DUT in the test. It could 37either be controlling a DUT set or an entire DUT pool. Implements the setup, 38cleanup and execute functions which spawn off multiple threads to 39control the execution of each step in the objects controlled. Should 40never be directly instantiated. 41 42CliqueDUTSet - Used to control a set within the DUT pool. It has a number of 43CliqueDUTRole objects to control. 44 45CliqueDUTPool - Used to control the entire DUT pool. It has a number of 46CliqueDUTSet objects to control. 47""" 48 49 50# Dummy result error reason to be used when exception is encountered in a role. 51ROLE_SETUP_EXCEPTION = "Role Setup Exception! " 52ROLE_EXECUTE_EXCEPTION = "Role Execute Exception! " 53ROLE_CLEANUP_EXCEPTION = "Role Teardown Exception! " 54 55# Dummy result error reason to be used when exception is encountered in a role. 56POOL_SETUP_EXCEPTION = "Pool Setup Exception! " 57POOL_CLEANUP_EXCEPTION = "Pool Teardown Exception! " 58 59# Result to returned after execution a sequence of steps. 60ControlResult = collections.namedtuple( 61 'ControlResult', [ 'uid', 'run_num', 'success', 62 'error_reason', 'start_time', 'end_time' ]) 63 64class CliqueDUTUnknownParamError(error.TestError): 65 """Indicates an error in finding a required param from the |test_params|.""" 66 pass 67 68 69class CliqueControl(object): 70 """CliqueControl is a base class which is used to control the DUT's in the 71 test. Not to be directly instantiated. 72 """ 73 74 def __init__(self, dut_objs, assoc_params=None, conn_worker=None, 75 test_params=None, uid=""): 76 """Initialize. 77 78 @param dut_objs: A list of objects that is being controlled by this 79 control object. 80 @param assoc_params: Association paramters to be used for this control 81 object. 82 @param conn_worker: ConnectionWorkerAbstract object, to run extra 83 work after successful connection. 84 @param test_params: A dictionary of params to be used for executing the 85 test. 86 @param uid: UID of this instance of the object. Host name for DUTRole 87 objects, Instance name for DUTBatch objects. 88 """ 89 self._dut_objs = dut_objs 90 self._test_params = test_params 91 self._assoc_params = assoc_params 92 self._conn_worker = conn_worker 93 self._uid = uid 94 95 def find_param(self, param_key): 96 """Find the relevant param value for a role from internal dictionary. 97 98 @param param_key: Look for the value of param_key in the dict. 99 100 @raises CliqueDUTUnknownParamError if there is an error in lookup. 101 """ 102 if not self._test_params.has_key(param_key): 103 raise CliqueDUTUnknownParamError("Param %s not found in %s" % 104 (param_key, self._test_params)) 105 return self._test_params.get(param_key) 106 107 @property 108 def dut_objs(self): 109 """Returns the dut_objs controlled by the object.""" 110 return self._dut_objs 111 112 @property 113 def dut_obj(self): 114 """Returns the first dut_obj controlled by the object.""" 115 return self._dut_objs[0] 116 117 @property 118 def uid(self): 119 """Returns a unique identifier associated with this object. It could 120 be just the hostname of the DUT in DUTRole objects or 121 set-number/pool-number in DUTSet DUTPool objects. 122 """ 123 return self._uid 124 125 @property 126 def assoc_params(self): 127 """Returns the association params corresponding to the object.""" 128 return self._assoc_params 129 130 @property 131 def conn_worker(self): 132 """Returns the connection worker corresponding to the object.""" 133 return self._conn_worker 134 135 136 def setup(self, run_num): 137 """Setup the DUT/DUT-set in the correct state before the sequence of 138 actions to be taken for the role is executed. 139 140 @param run_num: Run number of this execution. 141 142 @returns: An instance of ControlResult corresponding to all the errors 143 that were returned by the DUT/DUT's in the DUT-set which 144 is being controlled. 145 """ 146 pass 147 148 def cleanup(self, run_num): 149 """Cleanup the DUT/DUT-set state after the sequence of actions to be 150 taken for the role is executed. 151 152 @param run_num: Run number of this execution. 153 154 @returns: An instance of ControlResult corresponding to all the errors 155 that were returned by the DUT/DUT's in the DUT-set which 156 is being controlled. 157 """ 158 pass 159 160 def execute(self, run_num): 161 """Execute the sequence of actions to be taken for the role on the DUT 162 /DUT-set. 163 164 @param run_num: Run number of this execution. 165 166 @returns: An instance of ControlResult corresponding to all the errors 167 that were returned by the DUT/DUT's in the DUT-set which 168 is being controlled. 169 170 """ 171 pass 172 173 174class CliqueDUTRole(CliqueControl): 175 """CliqueDUTRole is a base class which defines the role entrusted to each 176 DUT in the Clique Test. Not to be directly instantiated. 177 """ 178 179 def __init__(self, dut, assoc_params=None, conn_worker=None, 180 test_params=None): 181 """Initialize. 182 183 @param dut: A DUTObject representing a DUT in the set. 184 @param assoc_params: Association paramters to be used for this role. 185 @param conn_worker: ConnectionWorkerAbstract object, to run extra 186 work after successful connection. 187 @param test_params: A dictionary of params to be used for executing the 188 test. 189 """ 190 super(CliqueDUTRole, self).__init__( 191 dut_objs=[dut], assoc_params=assoc_params, 192 conn_worker=conn_worker, test_params=test_params, 193 uid=dut.host.hostname) 194 195 def setup(self, run_num): 196 try: 197 assoc_params = self.assoc_params 198 self.dut_obj.wifi_client.shill.disconnect(assoc_params.ssid) 199 if not self.dut_obj.wifi_client.shill.init_test_network_state(): 200 result = ControlResult(uid=self.uid, 201 run_num=run_num, 202 success=False, 203 error_reason="Failed to set up isolated " 204 "test context profile.", 205 start_time="", 206 end_time="") 207 return result 208 else: 209 return None 210 except Exception as e: 211 result = ControlResult(uid=self.uid, 212 run_num=run_num, 213 success=False, 214 error_reason=ROLE_SETUP_EXCEPTION + str(e), 215 start_time="", 216 end_time="") 217 return result 218 219 def cleanup(self, run_num): 220 try: 221 self.dut_obj.wifi_client.shill.clean_profiles() 222 return None 223 except Exception as e: 224 result = ControlResult(uid=self.uid, 225 run_num=run_num, 226 success=False, 227 error_reason=ROLE_CLEANUP_EXCEPTION + str(e), 228 start_time="", 229 end_time="") 230 return result 231 232 def _connect_wifi(self, run_num): 233 """Helper function to make a connection to the associated AP.""" 234 assoc_params = self.assoc_params 235 logging.info('Connection attempt %d', run_num) 236 self.dut_obj.host.syslog('Connection attempt %d' % run_num) 237 start_time = self.dut_obj.host.run("date '+%FT%T.%N%:z'").stdout 238 start_time = start_time.strip() 239 assoc_result = xmlrpc_datatypes.deserialize( 240 self.dut_obj.wifi_client.shill.connect_wifi(assoc_params)) 241 end_time = self.dut_obj.host.run("date '+%FT%T.%N%:z'").stdout 242 end_time = end_time.strip() 243 success = assoc_result.success 244 if not success: 245 logging.error('Connection attempt %d failed; reason: %s', 246 run_num, assoc_result.failure_reason) 247 result = ControlResult(uid=self.uid, 248 run_num=run_num, 249 success=success, 250 error_reason=assoc_result.failure_reason, 251 start_time=start_time, 252 end_time=end_time) 253 return result 254 else: 255 logging.info('Connection attempt %d passed', run_num) 256 return None 257 258 def _disconnect_wifi(self): 259 """Helper function to disconnect from the associated AP.""" 260 assoc_params = self.assoc_params 261 self.dut_obj.wifi_client.shill.disconnect(assoc_params.ssid) 262 263 264# todo(rpius): Move these role implementations to a separate file since we'll 265# end up having a lot of roles defined. 266class DUTRoleConnectDisconnect(CliqueDUTRole): 267 """DUTRoleConnectDisconnect is used to make a DUT connect and disconnect 268 to a given AP repeatedly. 269 """ 270 271 def execute(self, run_num): 272 try: 273 result = self._connect_wifi(run_num) 274 if result: 275 return result 276 277 # Now disconnect from the AP. 278 self._disconnect_wifi() 279 280 return None 281 except Exception as e: 282 result = ControlResult(uid=self.uid, 283 run_num=run_num, 284 success=False, 285 error_reason=ROLE_EXECUTE_EXCEPTION + str(e), 286 start_time="", 287 end_time="") 288 return result 289 290 291class DUTRoleConnectDuration(CliqueDUTRole): 292 """DUTRoleConnectDuration is used to make a DUT connect to a given AP and 293 then check the liveness of the connection from another worker device. 294 """ 295 296 def setup(self, run_num): 297 result = super(DUTRoleConnectDuration, self).setup(run_num) 298 if result: 299 return result 300 # Let's check for the worker client now. 301 if not self.conn_worker: 302 return ControlResult(uid=self.uid, 303 run_num=run_num, 304 success=False, 305 error_reason="No connection worker found", 306 start_time="", 307 end_time="") 308 309 def execute(self, run_num): 310 try: 311 result = self._connect_wifi(run_num) 312 if result: 313 return result 314 315 # Let's start the ping from the worker client. 316 worker = connection_worker.ConnectionDuration.create_from_parent( 317 self.conn_worker) 318 worker.run(self.dut_obj.wifi_client) 319 320 return None 321 except Exception as e: 322 result = ControlResult(uid=self.uid, 323 run_num=run_num, 324 success=False, 325 error_reason=ROLE_EXECUTE_EXCEPTION + str(e), 326 start_time="", 327 end_time="") 328 return result 329 330 331def dut_batch_worker(dut_control_obj, method, error_results_queue, run_num): 332 """The method called by multiprocessing worker pool for running the DUT 333 control object's setup/execute/cleanup methods. This function is the 334 function which is repeatedly scheduled for each DUT/DUT-set through the 335 multiprocessing worker. This has to be defined outside the class because it 336 needs to be pickleable. 337 338 @param dut_control_obj: An object corresponding to DUT/DUT-set to control. 339 @param method: Method name to be invoked on the dut_control_obj. 340 it has to be one of setup/execute/teardown. 341 @param error_results_queue: Queue to put the error results after test. 342 @param run_num: Run number of this execution. 343 """ 344 logging.info("%s: Running %s", dut_control_obj.uid, method) 345 run_method = getattr(dut_control_obj, method, None) 346 if callable(run_method): 347 result = run_method(run_num) 348 if result: 349 error_results_queue.put(result) 350 351 352class CliqueDUTBatch(CliqueControl): 353 """CliqueDUTBatch is a base class which is used to control a batch of DUTs. 354 This could either be a DUT set or the entire DUT pool. Not to be directly 355 instantiated. 356 """ 357 # Used to store the instance number of derived classes. 358 BATCH_UID_NUM = {} 359 360 def __init__(self, dut_objs, test_params=None): 361 """Initialize. 362 363 @param dut_objs: A list of DUTRole objects representing the DUTs in set. 364 @param test_params: A dictionary of params to be used for executing the 365 test. 366 """ 367 uid_num = self.BATCH_UID_NUM.get(self.__class__.__name__, 1) 368 uid = self.__class__.__name__ + str(uid_num) 369 self.BATCH_UID_NUM[self.__class__.__name__] = uid_num + 1 370 super(CliqueDUTBatch, self).__init__( 371 dut_objs=dut_objs, test_params=test_params, uid=uid) 372 373 def _spawn_worker_threads(self, method, run_num): 374 """Spawns multiple threads to run the the |method(run_num)| on all the 375 control objects in parallel. 376 377 @param method: Method to be invoked on the dut_objs. 378 @param run_num: Run number of this execution. 379 380 @returns: An instance of ControlResult corresponding to all the errors 381 that were returned by the DUT/DUT's in the DUT-set which 382 is being controlled. 383 """ 384 tasks = [] 385 error_results_queue = multiprocessing.Queue() 386 for dut_obj in self.dut_objs: 387 task = multiprocessing.Process( 388 target=dut_batch_worker, 389 args=(dut_obj, method, error_results_queue, run_num)) 390 tasks.append(task) 391 # Run the tasks in parallel. 392 for task in tasks: 393 task.start() 394 for task in tasks: 395 task.join() 396 error_results = [] 397 while not error_results_queue.empty(): 398 result = error_results_queue.get() 399 # error_results returned at the DUT set level will be a list of 400 # ControlResult objects from each of the DUTs in the set. 401 # error_results returned at the DUT pool level will be a list of 402 # lists from each DUT set. Let's flatten out the list in that case 403 # since there could be ControlResult objects that are generated at 404 # the pool or set level which will make the final error result list 405 # assymetric where some elements are lists of ControlResult objects 406 # and some are just ControlResult objects. 407 if isinstance(result, list): 408 error_results.extend(result) 409 else: 410 error_results.append(result) 411 return error_results 412 413 def setup(self, run_num): 414 """Setup the DUT-set/pool in the correct state before the sequence of 415 actions to be taken for the role is executed. 416 417 @param run_num: Run number of this execution. 418 419 @returns: An instance of ControlResult corresponding to all the errors 420 that were returned by the DUT/DUT's in the DUT-set which 421 is being controlled. 422 """ 423 return self._spawn_worker_threads("setup", run_num) 424 425 def cleanup(self, run_num): 426 """Cleanup the DUT-set/pool state after the sequence of actions to be 427 taken for the role is executed. 428 429 @param run_num: Run number of this execution. 430 431 @returns: An instance of ControlResult corresponding to all the errors 432 that were returned by the DUT/DUT's in the DUT-set which 433 is being controlled. 434 """ 435 return self._spawn_worker_threads("cleanup", run_num) 436 437 def execute(self, run_num): 438 """Execute the sequence of actions to be taken for the role on the 439 DUT-set/pool. 440 441 @param run_num: Run number of this execution. 442 443 @returns: An instance of ControlResult corresponding to all the errors 444 that were returned by the DUT/DUT's in the DUT-set which 445 is being controlled. 446 447 """ 448 return self._spawn_worker_threads("execute", run_num) 449 450 451class CliqueDUTSet(CliqueDUTBatch): 452 """CliqueDUTSet is an object which is used to control all the DUT's in a DUT 453 set. 454 """ 455 def setup(self, run_num): 456 # Placeholder to add any set specific actions. 457 return super(CliqueDUTSet, self).setup(run_num) 458 459 def cleanup(self, run_num): 460 # Placeholder to add any set specific actions. 461 return super(CliqueDUTSet, self).cleanup(run_num) 462 463 def execute(self, run_num): 464 # Placeholder to add any set specific actions. 465 return super(CliqueDUTSet, self).execute(run_num) 466 467 468class CliqueDUTPool(CliqueDUTBatch): 469 """CliqueDUTSet is an object which is used to control all the DUT-sets in a 470 DUT pool. 471 """ 472 473 def setup(self, run_num): 474 # Let's start the packet capture before we kick off the entire pool 475 # execution. 476 try: 477 capturer = self.find_param('capturer') 478 capturer_frequency = self.find_param('capturer_frequency') 479 capturer_ht_type = self.find_param('capturer_ht_type') 480 capturer.start_capture(capturer_frequency, 481 width_type=capturer_ht_type) 482 except Exception as e: 483 result = ControlResult(uid=self.uid, 484 run_num=run_num, 485 success=False, 486 error_reason=POOL_SETUP_EXCEPTION + str(e), 487 start_time="", 488 end_time="") 489 # We cannot proceed with the test if this failed. 490 return result 491 # Now execute the setup on all the DUT-sets. 492 return super(CliqueDUTPool, self).setup(run_num) 493 494 def cleanup(self, run_num): 495 # First execute the cleanup on all the DUT-sets. 496 results = super(CliqueDUTPool, self).cleanup(run_num) 497 # Now stop the packet capture. 498 try: 499 capturer = self.find_param('capturer') 500 filename = str('connect_try_%d.trc' % (run_num)), 501 capturer.stop_capture(save_dir=self.outputdir, 502 save_filename=filename) 503 except Exception as e: 504 result = ControlResult(uid=self.uid, 505 run_num=run_num, 506 success=False, 507 error_reason=POOL_CLEANUP_EXCEPTION + str(e), 508 start_time="", 509 end_time="") 510 if results: 511 results.append(result) 512 else: 513 results = result 514 return results 515 516 def execute(self, run_num): 517 # Placeholder to add any pool specific actions. 518 return super(CliqueDUTPool, self).execute(run_num) 519 520 521def execute_dut_pool(dut_pool, dut_role_classes, assoc_params_list, 522 conn_workers, test_params, num_runs=1): 523 524 """Controls the DUT's in a given test scenario. The DUT's are assigned a 525 role according to the dut_role_classes provided for each DUT-set and all of 526 the sequence of steps are executed parallely on all the DUT's in the pool. 527 528 @param dut_pool: 2D list of DUT objects corresponding to the DUT's in the 529 DUT pool. 530 @param dut_role_classes: List of roles to be assigned to each set in the DUT 531 pool. Each element has to be a derived class of 532 CliqueDUTRole. 533 @param assoc_params_list: List of association parameters corrresponding 534 to the AP to test against for each set in the 535 DUT. 536 @param conn_workers: List of ConnectionWorkerAbstract objects, to 537 run extra work after successful connection. 538 @param test_params: List of params to be used for the test. 539 @num_runs: Number of iterations of the test to be run. 540 """ 541 # Every DUT set in the pool needs to have a corresponding DUT role, 542 # association parameters and connection worker assigned from the test. 543 # It is the responsibilty of the test scenario to make sure that there is a 544 # one to one mapping of all these elements since DUT control is going to 545 # be generic. 546 # This might mean that the test needs to duplicate the association 547 # parameters in the list if there is only 1 AP and 2 DUT sets. 548 # Or if there is no connection worker required, then the test should create 549 # a list of 'None' objects with length of 2. 550 # DUT control does not care if the same AP is used for 2 DUT sets or if the 551 # same connection worker is shared across 2 DUT sets as long as the 552 # length of the lists are equal. 553 554 if ((len(dut_pool) != len(dut_role_classes)) or 555 (len(dut_pool) != len(assoc_params_list)) or 556 (len(dut_pool) != len(conn_workers))): 557 raise error.TestError("Incorrect test configuration. Num DUT sets: %d, " 558 "Num DUT roles: %d, Num association params: %d, " 559 "Num connection workers: %d" % 560 (len(dut_pool), len(dut_role_classes), 561 len(assoc_params_list), len(conn_workers))) 562 563 dut_set_control_objs = [] 564 for dut_set, dut_role_class, assoc_params, conn_worker in \ 565 zip(dut_pool, dut_role_classes, assoc_params_list, conn_workers): 566 dut_control_objs = [] 567 for dut in dut_set: 568 dut_control_obj = dut_role_class( 569 dut, assoc_params, conn_worker, test_params) 570 dut_control_objs.append(dut_control_obj) 571 dut_set_control_obj = CliqueDUTSet(dut_control_objs, test_params) 572 dut_set_control_objs.append(dut_set_control_obj) 573 dut_pool_control_obj = CliqueDUTPool(dut_set_control_objs, test_params) 574 575 for run_num in range(0, num_runs): 576 # This setup, execute, cleanup calls on pool object, results in parallel 577 # invocation of call on all the DUT-sets which in turn results in 578 # parallel invocation of call on all the DUTs. 579 error_results = dut_pool_control_obj.setup(run_num) 580 if error_results: 581 return error_results 582 583 error_results = dut_pool_control_obj.execute(run_num) 584 if error_results: 585 # Try to cleanup before we leave. 586 dut_pool_control_obj.cleanup(run_num) 587 return error_results 588 589 error_results = dut_pool_control_obj.cleanup(run_num) 590 if error_results: 591 return error_results 592 return None 593