Support for extensions in chrome.py

Create a Chrome wrapper object.
Add support for loading component extensions by path.
Add a test invoking a private api in telemetry_LoginTest.

BUG=None

TEST=manual

Change-Id: Ice05decb47aeff5bebd445b33306cfc614decd0c
Reviewed-on: https://gerrit.chromium.org/gerrit/62259
Reviewed-by: Simran Basi <sbasi@chromium.org>
Commit-Queue: Achuith Bhandarkar <achuith@chromium.org>
Tested-by: Achuith Bhandarkar <achuith@chromium.org>
diff --git a/client/common_lib/cros/chrome.py b/client/common_lib/cros/chrome.py
index c353ed0..c7a37af 100644
--- a/client/common_lib/cros/chrome.py
+++ b/client/common_lib/cros/chrome.py
@@ -2,43 +2,61 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
-"""Telemetry-based Chrome automation functions."""
-
-from telemetry.core import browser_finder, browser_options
+from telemetry.core import browser_finder, browser_options, extension_to_load
 
 
-# The following constant is the name of the logged-in user that is specified
-# by the Telemetry chromeOS login extension
-# (chromium/src/tools/telemetry/telemetry/core/chrome/chromeos_login_ext).
-# The value here must match what is specified in that login extension.
+# Name of the logged-in user specified by the telemetry login extension.
 LOGIN_USER = 'test@test.test'
-_BROWSER_TYPE_LOGIN = 'system'
-_BROWSER_TYPE_GUEST = 'system-guest'
 
 
-def _get_browser(browser_type):
-    options = browser_options.BrowserOptions()
-    options.browser_type = browser_type
-    browser_to_create = browser_finder.FindBrowser(options)
-    return browser_to_create.Create()
+class Chrome(object):
+    """Wrapper for creating a telemetry browser instance with extensions."""
 
 
-def logged_in_browser():
-    """Returns a logged in browser.
-
-    Wrapping this within a Python with/as construct will take care of
-    automatically logging into Chrome at the start and logging out of Chrome
-    at the end, e.g.:
-
-    with chrome.logged_in_browser() as browser:
-        do_test()  # Will be logged in for this.
-    # Logged out at this point.
-
-    @return A Telemetry Browser object supporting context management.
-    """
-    return _get_browser(_BROWSER_TYPE_LOGIN)
+    BROWSER_TYPE_LOGIN = 'system'
+    BROWSER_TYPE_GUEST = 'system-guest'
 
 
-def incognito_browser():
-    """Returns an incognito browser."""
-    return _get_browser(_BROWSER_TYPE_GUEST)
+    def __init__(self, logged_in=True, extension_paths=[]):
+        options = browser_options.BrowserOptions()
+        self._browser_type = (self.BROWSER_TYPE_LOGIN
+                if logged_in else self.BROWSER_TYPE_GUEST)
+        options.browser_type = self._browser_type
+
+        if logged_in:
+            for path in extension_paths:
+                extension = extension_to_load.ExtensionToLoad(
+                        path, self._browser_type, is_component=True)
+                options.extensions_to_load.append(extension)
+            self._extensions_to_load = options.extensions_to_load
+
+        browser_to_create = browser_finder.FindBrowser(options)
+        self._browser = browser_to_create.Create()
+
+
+    def __enter__(self):
+        return self
+
+
+    def __exit__(self, *args):
+        self.browser.Close()
+
+
+    @property
+    def browser(self):
+        """Returns a telemetry browser instance."""
+        return self._browser
+
+
+    def get_extension(self, extension_path):
+        """Fetches a telemetry extension instance given the extension path."""
+        for ext in self._extensions_to_load:
+            if extension_path == ext.path:
+                return self.browser.extensions[ext]
+        return None
+
+
+    @property
+    def browser_type(self):
+        """Returns the browser_type."""
+        return self._browser_type