xref: /aosp_15_r20/external/autotest/client/common_lib/cros/chromedriver.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7)
1# Lint as: python2, python3
2# Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import atexit
7import logging
8import os
9from six.moves import urllib
10import six.moves.urllib.parse
11
12try:
13    from selenium import webdriver
14except ImportError:
15    # Ignore import error, as this can happen when builder tries to call the
16    # setup method of test that imports chromedriver.
17    logging.error('selenium module failed to be imported.')
18    pass
19
20from autotest_lib.client.bin import utils
21from autotest_lib.client.common_lib.cros import chrome
22
23CHROMEDRIVER_EXE_PATH = '/usr/local/chromedriver/chromedriver'
24X_SERVER_DISPLAY = ':0'
25X_AUTHORITY = '/home/chronos/.Xauthority'
26
27
28class chromedriver(object):
29    """Wrapper class, a context manager type, for tests to use Chrome Driver."""
30
31    def __init__(self,
32                 extra_chrome_flags=[],
33                 subtract_extra_chrome_flags=[],
34                 extension_paths=[],
35                 username=None,
36                 password=None,
37                 server_port=None,
38                 skip_cleanup=False,
39                 url_base=None,
40                 extra_chromedriver_args=None,
41                 gaia_login=False,
42                 disable_default_apps=True,
43                 dont_override_profile=False,
44                 chromeOptions={},
45                 *args,
46                 **kwargs):
47        """Initialize.
48
49        @param extra_chrome_flags: Extra chrome flags to pass to chrome, if any.
50        @param subtract_extra_chrome_flags: Remove default flags passed to
51                chrome by chromedriver, if any.
52        @param extension_paths: A list of paths to unzipped extensions. Note
53                                that paths to crx files won't work.
54        @param username: Log in using this username instead of the default.
55        @param password: Log in using this password instead of the default.
56        @param server_port: Port number for the chromedriver server. If None,
57                            an available port is chosen at random.
58        @param skip_cleanup: If True, leave the server and browser running
59                             so that remote tests can run after this script
60                             ends. Default is False.
61        @param url_base: Optional base url for chromedriver.
62        @param extra_chromedriver_args: List of extra arguments to forward to
63                                        the chromedriver binary, if any.
64        @param gaia_login: Logs in to real gaia.
65        @param disable_default_apps: For tests that exercise default apps.
66        @param dont_override_profile: Don't delete cryptohome before login.
67                                      Telemetry will output a warning with this
68                                      option.
69        """
70        if not isinstance(chromeOptions, dict):
71            raise TypeError("chromeOptions must be of type dict.")
72        self._cleanup = not skip_cleanup
73        assert os.geteuid() == 0, 'Need superuser privileges'
74
75        # When ChromeDriver starts Chrome on other platforms (Linux, Windows,
76        # etc.), it accepts flag inputs of the form "--flag_name" or
77        # "flag_name". Before starting Chrome with those flags, ChromeDriver
78        # reformats them all to "--flag_name". This behavior is copied
79        # to ChromeOS for consistency across platforms.
80        fixed_extra_chrome_flags = [
81            f if f.startswith('--') else '--%s' % f for f in extra_chrome_flags]
82
83        # Log in with telemetry
84        self._chrome = chrome.Chrome(extension_paths=extension_paths,
85                                     username=username,
86                                     password=password,
87                                     extra_browser_args=fixed_extra_chrome_flags,
88                                     gaia_login=gaia_login,
89                                     disable_default_apps=disable_default_apps,
90                                     dont_override_profile=dont_override_profile
91                                     )
92        self._browser = self._chrome.browser
93        # Close all tabs owned and opened by Telemetry, as these cannot be
94        # transferred to ChromeDriver.
95        self._browser.tabs[0].Close()
96
97        # Start ChromeDriver server
98        self._server = chromedriver_server(CHROMEDRIVER_EXE_PATH,
99                                           port=server_port,
100                                           skip_cleanup=skip_cleanup,
101                                           url_base=url_base,
102                                           extra_args=extra_chromedriver_args)
103
104        # Open a new tab using Chrome remote debugging. ChromeDriver expects
105        # a tab opened for remote to work. Tabs opened using Telemetry will be
106        # owned by Telemetry, and will be inaccessible to ChromeDriver.
107        urllib.request.urlopen('http://localhost:%i/json/new' %
108                        utils.get_chrome_remote_debugging_port())
109
110        chromeBaseOptions = {
111                'debuggerAddress':
112                ('localhost:%d' % utils.get_chrome_remote_debugging_port())
113        }
114        chromeOptions.update(chromeBaseOptions)
115        capabilities = {'chromeOptions':chromeOptions}
116        # Handle to chromedriver, for chrome automation.
117        try:
118            self.driver = webdriver.Remote(command_executor=self._server.url,
119                                           desired_capabilities=capabilities)
120        except NameError:
121            logging.error('selenium module failed to be imported.')
122            raise
123
124
125    def __enter__(self):
126        return self
127
128
129    def __exit__(self, *args):
130        """Clean up after running the test.
131
132        """
133        if hasattr(self, 'driver') and self.driver:
134            self.driver.close()
135            del self.driver
136
137        if not hasattr(self, '_cleanup') or self._cleanup:
138            if hasattr(self, '_server') and self._server:
139                self._server.close()
140                del self._server
141
142            if hasattr(self, '_browser') and self._browser:
143                self._browser.Close()
144                del self._browser
145
146    def get_extension(self, extension_path):
147        """Gets an extension by proxying to the browser.
148
149        @param extension_path: Path to the extension loaded in the browser.
150
151        @return: A telemetry extension object representing the extension.
152        """
153        return self._chrome.get_extension(extension_path)
154
155
156    @property
157    def chrome_instance(self):
158        """ The chrome instance used by this chrome driver instance. """
159        return self._chrome
160
161
162class chromedriver_server(object):
163    """A running ChromeDriver server.
164
165    This code is migrated from chrome:
166    src/chrome/test/chromedriver/server/server.py
167    """
168
169    def __init__(self, exe_path, port=None, skip_cleanup=False,
170                 url_base=None, extra_args=None):
171        """Starts the ChromeDriver server and waits for it to be ready.
172
173        Args:
174            exe_path: path to the ChromeDriver executable
175            port: server port. If None, an available port is chosen at random.
176            skip_cleanup: If True, leave the server running so that remote
177                          tests can run after this script ends. Default is
178                          False.
179            url_base: Optional base url for chromedriver.
180            extra_args: List of extra arguments to forward to the chromedriver
181                        binary, if any.
182        Raises:
183            RuntimeError if ChromeDriver fails to start
184        """
185        if not os.path.exists(exe_path):
186            raise RuntimeError('ChromeDriver exe not found at: ' + exe_path)
187
188        chromedriver_args = [exe_path]
189        if port:
190            # Allow remote connections if a port was specified
191            chromedriver_args.append('--whitelisted-ips')
192        else:
193            port = utils.get_unused_port()
194        chromedriver_args.append('--port=%d' % port)
195
196        self.url = 'http://localhost:%d' % port
197        if url_base:
198            chromedriver_args.append('--url-base=%s' % url_base)
199            self.url = six.moves.urllib.parse.urljoin(self.url, url_base)
200
201        if extra_args:
202            chromedriver_args.extend(extra_args)
203
204        # TODO(ihf): Remove references to X after M45.
205        # Chromedriver will look for an X server running on the display
206        # specified through the DISPLAY environment variable.
207        os.environ['DISPLAY'] = X_SERVER_DISPLAY
208        os.environ['XAUTHORITY'] = X_AUTHORITY
209
210        self.bg_job = utils.BgJob(chromedriver_args, stderr_level=logging.DEBUG)
211        if self.bg_job is None:
212            raise RuntimeError('ChromeDriver server cannot be started')
213
214        try:
215            timeout_msg = 'Timeout on waiting for ChromeDriver to start.'
216            utils.poll_for_condition(self.is_running,
217                                     exception=utils.TimeoutError(timeout_msg),
218                                     timeout=10,
219                                     sleep_interval=.1)
220        except utils.TimeoutError:
221            self.close_bgjob()
222            raise RuntimeError('ChromeDriver server did not start')
223
224        logging.debug('Chrome Driver server is up and listening at port %d.',
225                      port)
226        if not skip_cleanup:
227            atexit.register(self.close)
228
229
230    def is_running(self):
231        """Returns whether the server is up and running."""
232        try:
233            urllib.request.urlopen(self.url + '/status')
234            return True
235        except urllib.error.URLError as e:
236            return False
237
238
239    def close_bgjob(self):
240        """Close background job and log stdout and stderr."""
241        utils.nuke_subprocess(self.bg_job.sp)
242        utils.join_bg_jobs([self.bg_job], timeout=1)
243        result = self.bg_job.result
244        if result.stdout or result.stderr:
245            logging.info('stdout of Chrome Driver:\n%s', result.stdout)
246            logging.error('stderr of Chrome Driver:\n%s', result.stderr)
247
248
249    def close(self):
250        """Kills the ChromeDriver server, if it is running."""
251        if self.bg_job is None:
252            return
253
254        try:
255            urllib.request.urlopen(self.url + '/shutdown', timeout=10).close()
256        except:
257            pass
258
259        self.close_bgjob()
260