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