Haketilo

from selenium.webdriver.support.ui import WebDriverWait

"""

tab_id_responder is meant to be appended to background script of a test

extension.

"""

tab_id_responder = '''

function tell_tab_id(msg, sender, respond_cb) {

    if (msg[0] === "learn_tab_id")

        respond_cb(sender.tab.id);

}

browser.runtime.onMessage.addListener(tell_tab_id);

'''

"""

tab_id_asker is meant to be appended to content script of a test extension.

"""

tab_id_asker = '''

browser.runtime.sendMessage(["learn_tab_id"])

    .then(tid => window.wrappedJSObject.haketilo_tab = tid);

'''

def run_content_script_in_new_window(driver, url):

    """

    Expect an extension to be loaded which had tab_id_responder and tab_id_asker

    appended to its background and content scripts, respectively.

    Open the provided url in a new tab, find its tab id and return it, with

    current window changed back to the initial one.

    """

    initial_handle = driver.current_window_handle

    handles = driver.window_handles

    driver.execute_script('window.open(arguments[0], "_blank");', url)

    WebDriverWait(driver, 10).until(lambda d: d.window_handles is not handles)

    new_handle = [h for h in driver.window_handles if h not in handles][0]

    driver.switch_to.window(new_handle)

    get_tab_id = lambda d: d.execute_script('return window.haketilo_tab;')

    tab_id = WebDriverWait(driver, 10).until(get_tab_id)

    driver.switch_to.window(initial_handle)

    return tab_id