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