1# Lint as: python2, python3
2from __future__ import absolute_import
3from __future__ import division
4from __future__ import print_function
5
6import logging
7import os
8from six.moves import range
9import time
10
11from autotest_lib.client.common_lib import error
12from autotest_lib.client.common_lib import utils
13
14
15class UIScreenshoter(object):
16    """Simple class to take screenshots within the ui_utils framework."""
17
18    _SCREENSHOT_DIR_PATH = '/var/log/ui_utils'
19    _SCREENSHOT_BASENAME = 'ui-screenshot'
20
21    def __init__(self):
22        if not os.path.exists(self._SCREENSHOT_DIR_PATH):
23            os.mkdir(self._SCREENSHOT_DIR_PATH, 0o755)
24        self.screenshot_num = 0
25
26    def take_ss(self):
27        try:
28            utils.run('screenshot "{}/{}_iter{}.png"'.format(
29                      self._SCREENSHOT_DIR_PATH, self._SCREENSHOT_BASENAME,
30                      self.screenshot_num))
31            self.screenshot_num += 1
32        except Exception as e:
33            logging.warning('Unable to capture screenshot. %s', e)
34
35class UI_Handler(object):
36
37    REGEX_ALL = '/(.*?)/'
38
39    PROMISE_TEMPLATE = \
40        '''new Promise(function(resolve, reject) {
41            chrome.automation.getDesktop(function(root) {
42                    resolve(%s);
43                })
44            })'''
45
46    def __init__(self):
47        self.screenshoter = UIScreenshoter()
48
49    def start_ui_root(self, cr):
50        """Start the UI root object for testing."""
51        self.ext = cr.autotest_ext
52
53    def is_obj_restricted(self, name, isRegex=False, role=None):
54        """
55        Return True if the object restriction is 'disabled'.
56        This usually means the button is either greyed out, locked, etc.
57
58        @param name: Parameter to provide to the 'name' attribute.
59        @param isRegex: bool, if the item is a regex.
60        @param role: Parameter to provide to the 'role' attribute.
61
62        """
63        FindParams = self._get_FindParams_str(name=name,
64                                              role=role,
65                                              isRegex=isRegex)
66        try:
67            restriction = self.ext.EvaluateJavaScript(
68                self.PROMISE_TEMPLATE % ("%s.restriction" % FindParams),
69                promise=True)
70        except Exception:
71            raise error.TestError(
72                'Could not find object {}.'.format(name))
73        if restriction == 'disabled':
74            return True
75        return False
76
77    def item_present(self, name, isRegex=False, flip=False, role=None):
78        """
79        Determines if an object is present on the screen
80
81        @param name: Parameter to provide to the 'name' attribute.
82        @param isRegex: bool, if the 'name' is a regex.
83        @param flip: Flips the return status.
84        @param role: Parameter to provide to the 'role' attribute.
85
86        @returns:
87            True if object is present and flip is False.
88            False if object is present and flip is True.
89            False if object is not present and flip is False.
90            True if object is not present and flip is True.
91
92        """
93        FindParams = self._get_FindParams_str(name=name,
94                                              role=role,
95                                              isRegex=isRegex)
96
97        if flip is True:
98            return not self._is_item_present(FindParams)
99        return self._is_item_present(FindParams)
100
101    def wait_for_ui_obj(self,
102                        name,
103                        isRegex=False,
104                        remove=False,
105                        role=None,
106                        timeout=10):
107        """
108        Waits for the UI object specified.
109
110        @param name: Parameter to provide to the 'name' attribute.
111        @param isRegex: bool, if the 'name' is a regex.
112        @param remove: bool, if you are waiting for the item to be removed.
113        @param role: Parameter to provide to the 'role' attribute.
114        @param timeout: int, time to wait for the item.
115
116        @raises error.TestError if the element is not loaded (or removed).
117
118        """
119        try:
120            utils.poll_for_condition(
121                condition=lambda: self.item_present(name=name,
122                                                    isRegex=isRegex,
123                                                    flip=remove,
124                                                    role=role),
125                timeout=timeout,
126                exception=error.TestError('{} did not load'
127                                          .format(name)))
128        except error.TestError:
129            self.screenshoter.take_ss()
130            logging.debug("CURRENT UI ITEMS VISIBLE {}".format(
131                self.list_screen_items()))
132            raise error.TestError('{} did not load'.format(name))
133
134    def did_obj_not_load(self, name, isRegex=False, timeout=5):
135        """
136        Specifically used to wait and see if an item appears on the UI.
137
138        NOTE: This is different from wait_for_ui_obj because that returns as
139        soon as the object is either loaded or not loaded. This function will
140        wait to ensure over the timeout period the object never loads.
141        Additionally it will return as soon as it does load. Basically a fancy
142        time.sleep()
143
144        @param name: Parameter to provide to the 'name' attribute.
145        @param isRegex: bool, if the item is a regex.
146        @param timeout: Time in seconds to wait for the object to appear.
147
148        @returns: True if object never loaded within the timeout period,
149            else False.
150
151        """
152        t1 = time.time()
153        while time.time() - t1 < timeout:
154            if self.item_present(name=name, isRegex=isRegex):
155                return False
156            time.sleep(1)
157        return True
158
159    def doDefault_on_obj(self, name, isRegex=False, role=None):
160        """Runs the .doDefault() js command on the element."""
161        FindParams = self._get_FindParams_str(name=name,
162                                              role=role,
163                                              isRegex=isRegex)
164        try:
165            self.ext.EvaluateJavaScript(
166                self.PROMISE_TEMPLATE % ("%s.doDefault()" % FindParams),
167                promise=True)
168        except:
169            logging.info('Unable to .doDefault() on {}. All items: {}'
170                         .format(FindParams, self.list_screen_items()))
171            raise error.TestError("doDefault failed on {}".format(FindParams))
172
173    def doCommand_on_obj(self, name, cmd, isRegex=False, role=None):
174        """Run the specified command on the element."""
175        FindParams = self._get_FindParams_str(name=name,
176                                              role=role,
177                                              isRegex=isRegex)
178        return self.ext.EvaluateJavaScript(self.PROMISE_TEMPLATE % """
179            %s.%s""" % (FindParams, cmd), promise=True)
180
181    def list_screen_items(self,
182                          role=None,
183                          name=None,
184                          isRegex=False,
185                          attr='name'):
186
187        """
188        List all the items currently visable on the screen.
189
190        If no paramters are given, it will return the name of each item,
191        including items with empty names.
192
193        @param role: The role of the items to use (ie button).
194        @param name: Parameter to provide to the 'name' attribute.
195        @param isRegex: bool, if the obj is a regex.
196        @param attr: Str, the attribute you want returned in the list
197            (eg 'name').
198
199        """
200
201        if isRegex:
202            if name is None:
203                raise error.TestError('If regex is True name must be given')
204            name = self._format_obj(name, isRegex)
205        elif name is not None:
206            name = self._format_obj(name, isRegex)
207        name = self.REGEX_ALL if name is None else name
208        role = self.REGEX_ALL if role is None else self._format_obj(role,
209                                                                    False)
210
211        new_promise = self.PROMISE_TEMPLATE % """root.findAll({attributes:
212            {name: %s, role: %s}}).map(node => node.%s)""" % (name, role, attr)
213        return self.ext.EvaluateJavaScript(new_promise, promise=True)
214
215    def get_name_role_list(self, name=None, role=None):
216        """
217        Return [{}, {}] containing the name/role of everything on screen.
218
219        """
220        name = self.REGEX_ALL if name is None else name
221        role = self.REGEX_ALL if role is None else self._format_obj(role,
222                                                                    False)
223
224        new_promise = self.PROMISE_TEMPLATE % """root.findAll({attributes:
225            {name: %s, role: %s}}).map(node =>
226            {return {name: node.name, role: node.role} })""" % (name, role)
227
228        return self.ext.EvaluateJavaScript(new_promise, promise=True)
229
230    def _format_obj(self, name, isRegex):
231        """
232        Formats the object for use in the javascript name attribute.
233
234        When searching for an element on the UI, a regex expression or string
235        can be used. If the search is using a string, the obj will need to be
236        wrapped in quotes. A Regex is not.
237
238        @param name: Parameter to provide to the 'name' attribute.
239        @param isRegex: if True, the object will be returned as is, if False
240            the obj will be returned wrapped in quotes.
241
242        @returns: The formatted string for regex/name.
243        """
244        if isRegex:
245            return name
246        else:
247            return '"{}"'.format(name)
248
249    def _get_FindParams_str(self, name, role, isRegex):
250        """Returns the FindParms string, so that automation node functions
251        can be run on it
252
253        @param role: The role of the items to use (ie button).
254        @param name: Parameter to provide to the 'name' attribute.
255        @param isRegex: bool, if the obj is a regex.
256
257        @returns: The ".find($FindParams)" string, which can be used to run
258            automation node commands, such as .doDefault()
259
260        """
261        FINDPARAMS_BASE = """
262        root.find({attributes:
263                  {name: %s,
264                   role: %s}}
265                 )"""
266
267        name = self._format_obj(name, isRegex)
268        if role is None:
269            role = self.REGEX_ALL
270        else:
271            role = self._format_obj(role, False)
272        return (FINDPARAMS_BASE % (name, role))
273
274    def _is_item_present(self, findParams):
275        """Return False if tempVar is None, else True."""
276        item_present = self.ext.EvaluateJavaScript(
277            self.PROMISE_TEMPLATE % findParams,
278            promise=True)
279        if item_present is None:
280            return False
281        return True
282
283    def click_and_wait_for_item_with_retries(self,
284                                             item_to_click,
285                                             item_to_wait_for,
286                                             isRegex_click=False,
287                                             isRegex_wait=False,
288                                             click_role=None,
289                                             wait_role=None):
290        """
291        Click on an item, and wait for a subsequent item to load. If the new
292        item does not load, we attempt to click the button again.
293
294        This being done to remove the corner case of button being visually
295        loaded, but not fully ready to be clicked yet. In simple terms:
296            Click button --> Check for next button to appear
297            IF button does not appear, its likely the original button click did
298            not work, thus reclick that button --> recheck for the next button.
299
300            If button did appear, stop clicking.
301
302        """
303        self.doDefault_on_obj(item_to_click,
304                              role=click_role,
305                              isRegex=isRegex_click)
306        for retry in range(3):
307            try:
308                self.wait_for_ui_obj(item_to_wait_for,
309                                     role=wait_role,
310                                     isRegex=isRegex_wait,
311                                     timeout=6)
312                break
313            except error.TestError:
314                self.doDefault_on_obj(item_to_click,
315                                      role=click_role,
316                                      isRegex=isRegex_click)
317        else:
318            raise error.TestError('Item {} did not load after 2 tries'.format(
319                                  item_to_wait_for))
320