1# Copyright 2018 Google Inc.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Module for Mobly controller management."""
15import collections
16import copy
17import logging
18import yaml
19
20from mobly import expects
21from mobly import records
22from mobly import signals
23
24
25def verify_controller_module(module):
26  """Verifies a module object follows the required interface for
27  controllers.
28
29  The interface is explained in the docstring of
30  `base_test.BaseTestClass.register_controller`.
31
32  Args:
33    module: An object that is a controller module. This is usually
34      imported with import statements or loaded by importlib.
35
36  Raises:
37    ControllerError: if the module does not match the Mobly controller
38      interface, or one of the required members is null.
39  """
40  required_attributes = ('create', 'destroy', 'MOBLY_CONTROLLER_CONFIG_NAME')
41  for attr in required_attributes:
42    if not hasattr(module, attr):
43      raise signals.ControllerError(
44          'Module %s missing required controller module attribute %s.'
45          % (module.__name__, attr)
46      )
47    if not getattr(module, attr):
48      raise signals.ControllerError(
49          'Controller interface %s in %s cannot be null.'
50          % (attr, module.__name__)
51      )
52
53
54class ControllerManager:
55  """Manages the controller objects for Mobly tests.
56
57  This manages the life cycles and info retrieval of all controller objects
58  used in a test.
59
60  Attributes:
61    controller_configs: dict, controller configs provided by the user via
62      test bed config.
63  """
64
65  def __init__(self, class_name, controller_configs):
66    # Controller object management.
67    self._controller_objects = (
68        collections.OrderedDict()
69    )  # controller_name: objects
70    self._controller_modules = {}  # controller_name: module
71    self._class_name = class_name
72    self.controller_configs = controller_configs
73
74  def register_controller(self, module, required=True, min_number=1):
75    """Loads a controller module and returns its loaded devices.
76
77    This is to be used in a mobly test class.
78
79    Args:
80      module: A module that follows the controller module interface.
81      required: A bool. If True, failing to register the specified
82        controller module raises exceptions. If False, the objects
83        failed to instantiate will be skipped.
84      min_number: An integer that is the minimum number of controller
85        objects to be created. Default is one, since you should not
86        register a controller module without expecting at least one
87        object.
88
89    Returns:
90      A list of controller objects instantiated from controller_module, or
91      None if no config existed for this controller and it was not a
92      required controller.
93
94    Raises:
95      ControllerError:
96        * The controller module has already been registered.
97        * The actual number of objects instantiated is less than the
98        * `min_number`.
99        * `required` is True and no corresponding config can be found.
100        * Any other error occurred in the registration process.
101    """
102    verify_controller_module(module)
103    # Use the module's name as the ref name
104    module_ref_name = module.__name__.split('.')[-1]
105    if module_ref_name in self._controller_objects:
106      raise signals.ControllerError(
107          'Controller module %s has already been registered. It cannot '
108          'be registered again.' % module_ref_name
109      )
110    # Create controller objects.
111    module_config_name = module.MOBLY_CONTROLLER_CONFIG_NAME
112    if module_config_name not in self.controller_configs:
113      if required:
114        raise signals.ControllerError(
115            'No corresponding config found for %s' % module_config_name
116        )
117      logging.warning(
118          'No corresponding config found for optional controller %s',
119          module_config_name,
120      )
121      return None
122    try:
123      # Make a deep copy of the config to pass to the controller module,
124      # in case the controller module modifies the config internally.
125      original_config = self.controller_configs[module_config_name]
126      controller_config = copy.deepcopy(original_config)
127      objects = module.create(controller_config)
128    except Exception:
129      logging.exception(
130          'Failed to initialize objects for controller %s, abort!',
131          module_config_name,
132      )
133      raise
134    if not isinstance(objects, list):
135      raise signals.ControllerError(
136          'Controller module %s did not return a list of objects, abort.'
137          % module_ref_name
138      )
139    # Check we got enough controller objects to continue.
140    actual_number = len(objects)
141    if actual_number < min_number:
142      module.destroy(objects)
143      raise signals.ControllerError(
144          'Expected to get at least %d controller objects, got %d.'
145          % (min_number, actual_number)
146      )
147    # Save a shallow copy of the list for internal usage, so tests can't
148    # affect internal registry by manipulating the object list.
149    self._controller_objects[module_ref_name] = copy.copy(objects)
150    logging.debug(
151        'Found %d objects for controller %s', len(objects), module_config_name
152    )
153    self._controller_modules[module_ref_name] = module
154    return objects
155
156  def unregister_controllers(self):
157    """Destroy controller objects and clear internal registry.
158
159    This will be called after each test class.
160    """
161    # TODO(xpconanfan): actually record these errors instead of just
162    # logging them.
163    for name, module in self._controller_modules.items():
164      logging.debug('Destroying %s.', name)
165      with expects.expect_no_raises('Exception occurred destroying %s.' % name):
166        module.destroy(self._controller_objects[name])
167    self._controller_objects = collections.OrderedDict()
168    self._controller_modules = {}
169
170  def _create_controller_info_record(self, controller_module_name):
171    """Creates controller info record for a particular controller type.
172
173    Info is retrieved from all the controller objects spawned from the
174    specified module, using the controller module's `get_info` function.
175
176    Args:
177      controller_module_name: string, the name of the controller module
178        to retrieve info from.
179
180    Returns:
181      A records.ControllerInfoRecord object.
182    """
183    module = self._controller_modules[controller_module_name]
184    controller_info = None
185    try:
186      controller_info = module.get_info(
187          copy.copy(self._controller_objects[controller_module_name])
188      )
189    except AttributeError:
190      logging.warning(
191          'No optional debug info found for controller '
192          '%s. To provide it, implement `get_info`.',
193          controller_module_name,
194      )
195    try:
196      yaml.dump(controller_info)
197    except TypeError:
198      logging.warning(
199          'The info of controller %s in class "%s" is not '
200          'YAML serializable! Coercing it to string.',
201          controller_module_name,
202          self._class_name,
203      )
204      controller_info = str(controller_info)
205    return records.ControllerInfoRecord(
206        self._class_name, module.MOBLY_CONTROLLER_CONFIG_NAME, controller_info
207    )
208
209  def get_controller_info_records(self):
210    """Get the info records for all the controller objects in the manager.
211
212    New info records for each controller object are created for every call
213    so the latest info is included.
214
215    Returns:
216      List of records.ControllerInfoRecord objects. Each opject conatins
217      the info of a type of controller
218    """
219    info_records = []
220    for controller_module_name in self._controller_objects.keys():
221      with expects.expect_no_raises(
222          'Failed to collect controller info from %s' % controller_module_name
223      ):
224        record = self._create_controller_info_record(controller_module_name)
225        if record:
226          info_records.append(record)
227    return info_records
228