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