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