facilitate broadcasting messages to different execution contexts within the webextension

6 files changed, 229 insertions, 21 deletions

diff --git a/test/unit/conftest.py b/test/unit/conftest.py
index e1c98a1..eec311c 100644
--- a/test/unit/conftest.py
+++ b/test/unit/conftest.py
@@ -27,6 +27,7 @@ Common fixtures for Haketilo unit tests
 
 import pytest
 from pathlib import Path
+from selenium.webdriver.common.by import By
 from selenium.webdriver.support.ui import WebDriverWait
 from selenium.webdriver.support import expected_conditions as EC
 
@@ -42,11 +43,24 @@ def proxy():
     httpd.shutdown()
 
 @pytest.fixture(scope="package")
-def driver(proxy):
+def _driver(proxy):
     with firefox_safe_mode() as driver:
         yield driver
         driver.quit()
 
+def close_all_but_one_window(driver):
+    while len(driver.window_handles) > 1:
+        driver.switch_to.window(driver.window_handles[-1])
+        driver.close()
+    driver.switch_to.window(driver.window_handles[0])
+
+@pytest.fixture()
+def driver(_driver, request):
+    nav_target = request.node.get_closest_marker('get_page')
+    close_all_but_one_window(_driver)
+    _driver.get(nav_target.args[0] if nav_target else 'about:blank')
+    yield _driver
+
 @pytest.fixture()
 def webextension(driver, request):
     ext_data = request.node.get_closest_marker('ext_data')
@@ -58,7 +72,7 @@ def webextension(driver, request):
     driver.get('https://gotmyowndoma.in/')
     addon_id = driver.install_addon(str(ext_path), temporary=True)
     WebDriverWait(driver, 10).until(
-        EC.title_contains("Extension's options page for testing")
+        EC.url_matches('^moz-extension://.*')
     )
     yield
     driver.uninstall_addon(addon_id)
@@ -115,22 +129,28 @@ def _execute_in_page_context(driver, script, args):
 
         raise e from None
 
-@pytest.fixture(scope="package")
-def execute_in_page(driver):
-    def do_execute(script, *args, **kwargs):
-        if 'page' in kwargs:
-            driver.get(kwargs['page'])
+# Some fixtures here just define functions that operate on driver. We should
+# consider making them into webdriver wrapper class methods.
 
+@pytest.fixture()
+def execute_in_page(driver):
+    def do_execute(script, *args):
         return _execute_in_page_context(driver, script, args)
 
     yield do_execute
 
-@pytest.fixture(scope="package")
+@pytest.fixture()
 def load_into_page(driver):
-    def do_load(path, import_dirs, *args, **kwargs):
-        if 'page' in kwargs:
-            driver.get(kwargs['page'])
-
+    def do_load(path, import_dirs, *args):
         _execute_in_page_context(driver, load_script(path, import_dirs), args)
 
     yield do_load
+
+@pytest.fixture()
+def wait_elem_text(driver):
+    def do_wait(id, text):
+        WebDriverWait(driver, 10).until(
+            EC.text_to_be_present_in_element((By.ID, id), text)
+        )
+
+    yield do_wait
diff --git a/test/unit/test_basic.py b/test/unit/test_basic.py
index 3b09cb6..ca956e7 100644
--- a/test/unit/test_basic.py
+++ b/test/unit/test_basic.py
@@ -31,18 +31,19 @@ def test_driver(driver):
         )
         assert "Schrodinger's Document" in title
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_script_loader(execute_in_page, load_into_page):
     """
     A trivial test case that verifies Haketilo's .js files can be properly
     loaded into a test page together with their dependencies.
     """
-    load_into_page('common/stored_types.js', ['common'],
-                   page='https://gotmyowndoma.in')
+    load_into_page('common/stored_types.js', ['common'])
 
     assert execute_in_page('returnval(TYPE_PREFIX.VAR);') == '_'
 
 @pytest.mark.ext_data({})
-def test_webextension(driver, webextension):
+@pytest.mark.usefixtures('webextension')
+def test_webextension(driver):
     """
     A trivial test case that verifies a test WebExtension created and installed
     by the `webextension` fixture works and redirects specially-constructed URLs
diff --git a/test/unit/test_broadcast.py b/test/unit/test_broadcast.py
new file mode 100644
index 0000000..c8c19d1
--- /dev/null
+++ b/test/unit/test_broadcast.py
@@ -0,0 +1,181 @@
+# SPDX-License-Identifier: CC0-1.0
+
+"""
+Haketilo unit tests - message broadcasting
+"""
+
+# This file is part of Haketilo
+#
+# Copyright (C) 2021, Wojtek Kosior
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the CC0 1.0 Universal License as published by
+# the Creative Commons Corporation.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# CC0 1.0 Universal License for more details.
+
+import pytest
+
+from ..script_loader import load_script
+
+def broker_js():
+    return load_script('background/broadcast_broker.js',
+                       ['common', 'background']) + ';start_broadcast_broker();'
+
+def broadcast_js():
+    return load_script('common/broadcast.js', ['common'])
+
+test_page_html = '''
+<!DOCTYPE html>
+<script src="/testpage.js"></script>
+<h2>d0 (channel `somebodyoncetoldme`)</h2>
+<div id="d0"></div>
+<h2>d1 (channel `worldisgonnarollme`)</h2>
+<div id="d1"></div>
+<h2>d2 (both channels)</h2>
+<div id="d2"></div>
+'''
+
+@pytest.mark.ext_data({
+    'background_script': broker_js,
+    'test_page':         test_page_html,
+    'extra_files': {
+        'testpage.js':   broadcast_js
+    }
+})
+@pytest.mark.usefixtures('webextension')
+def test_broadcast(driver, execute_in_page, wait_elem_text):
+    """
+    A test that verifies the broadcasting system based on WebExtension messaging
+    API and implemented in `background/broadcast_broker.js` and
+    `common/broadcast.js` works correctly.
+    """
+
+    # The broadcast facility is meant to enable message distribution between
+    # multiple contexts (e.g. different tabs/windows). Let's open the same
+    # extension's test page in a second window.
+    driver.execute_script(
+        '''
+        window.open(window.location.href, "_blank");
+        window.open(window.location.href, "_blank");
+        ''')
+    windows = [*driver.window_handles]
+    assert len(windows) == 3
+
+    # Let's first test if a simple message can be successfully broadcasted
+    driver.switch_to.window(windows[0])
+    execute_in_page(
+        '''
+        const divs = [0, 1, 2].map(n => document.getElementById("d" + n));
+        let appender = n => (t => divs[n].append("\\n" + `[${t[0]}, ${t[1]}]`));
+        let listener0 = broadcast.listener_connection(appender(0));
+        broadcast.subscribe(listener0, "somebodyoncetoldme");
+        ''')
+
+    driver.switch_to.window(windows[1])
+    execute_in_page(
+        '''
+        let sender0 = broadcast.sender_connection();
+        broadcast.out(sender0, "somebodyoncetoldme", "iaintthesharpesttool");
+        ''')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d0', '[somebodyoncetoldme, iaintthesharpesttool]')
+
+    # Let's add 2 more listeners
+    driver.switch_to.window(windows[0])
+    execute_in_page(
+        '''
+        let listener1 = broadcast.listener_connection(appender(1));
+        broadcast.subscribe(listener1, "worldisgonnarollme");
+        let listener2 = broadcast.listener_connection(appender(2));
+        broadcast.subscribe(listener2, "worldisgonnarollme");
+        broadcast.subscribe(listener2, "somebodyoncetoldme");
+        ''')
+
+    # Let's send one message to one channel and one to the other. Verify they
+    # were received by the rght listeners.
+    driver.switch_to.window(windows[1])
+    execute_in_page(
+        '''
+        broadcast.out(sender0, "somebodyoncetoldme", "intheshed");
+        broadcast.out(sender0, "worldisgonnarollme", "shewaslooking");
+        ''')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d0', 'intheshed')
+    wait_elem_text('d1', 'shewaslooking')
+    wait_elem_text('d2', 'intheshed')
+    wait_elem_text('d2', 'shewaslooking')
+
+    text = execute_in_page('returnval(divs[0].innerText);')
+    assert 'shewaslooking' not in text
+    text = execute_in_page('returnval(divs[1].innerText);')
+    assert 'intheshed' not in text
+
+    # Let's create a second sender in third window and use it to send messages
+    # with the 'prepare' feature.
+    driver.switch_to.window(windows[2])
+    execute_in_page(
+        '''
+        let sender1 = broadcast.sender_connection();
+        broadcast.prepare(sender1, "somebodyoncetoldme", "kindadumb");
+        broadcast.out(sender1, "worldisgonnarollme", "withherfinger");
+        ''')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d1', 'withherfinger')
+    text = execute_in_page('returnval(divs[0].innerText);')
+    assert 'kindadumb' not in text
+
+    driver.switch_to.window(windows[2])
+    execute_in_page('broadcast.flush(sender1);')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d0', 'kindadumb')
+
+    # Let's verify that prepare()'d messages are properly discarded when
+    # discard() is called.
+    driver.switch_to.window(windows[2])
+    execute_in_page(
+        '''
+        broadcast.prepare(sender1, "somebodyoncetoldme", "andherthumb");
+        broadcast.discard(sender1);
+        broadcast.prepare(sender1, "somebodyoncetoldme", "andhermiddlefinger");
+        broadcast.flush(sender1);
+        ''')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d0', 'andhermiddlefinger')
+    text = execute_in_page('returnval(divs[0].innerText);')
+    assert 'andherthumb' not in text
+
+    # Let's verify prepare()'d messages are properly auto-flushed when the other
+    # end of the connection gets killed (e.g. because browser tab gets closed).
+    driver.switch_to.window(windows[2])
+    execute_in_page(
+        '''
+        broadcast.prepare(sender1, "worldisgonnarollme", "intheshape", 500);
+        ''')
+    driver.close()
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d2', 'intheshape')
+
+    # Verify listener's connection gets closed properly.
+    execute_in_page('broadcast.close(listener0); broadcast.close(listener1);')
+
+    driver.switch_to.window(windows[1])
+    execute_in_page('broadcast.out(sender0, "worldisgonnarollme", "ofanL");')
+    execute_in_page('broadcast.out(sender0, "somebodyoncetoldme", "forehead");')
+
+    driver.switch_to.window(windows[0])
+    wait_elem_text('d2', 'ofanL')
+    wait_elem_text('d2', 'forehead')
+    for i in (0, 1):
+        text = execute_in_page('returnval(divs[arguments[0]].innerText);', i)
+        assert 'ofanL' not in text
+        assert 'forehead' not in text
diff --git a/test/unit/test_indexeddb.py b/test/unit/test_indexeddb.py
index f1322fb..965f318 100644
--- a/test/unit/test_indexeddb.py
+++ b/test/unit/test_indexeddb.py
@@ -47,12 +47,13 @@ sample_files_by_hash = dict([[file['hash_key'], file['contents']]
 def file_ref(file_name):
     return {'file': file_name, 'hash_key': sample_files[file_name]['hash_key']}
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_save_remove_item(execute_in_page, indexeddb_code):
     """
     indexeddb.js facilitates operating on Haketilo's internal database.
     Verify database operations work properly.
     """
-    execute_in_page(indexeddb_code, page='https://gotmyowndoma.in')
+    execute_in_page(indexeddb_code)
     # Don't use Haketilo's default initial data.
     execute_in_page(
         '''{
diff --git a/test/unit/test_patterns.py b/test/unit/test_patterns.py
index 802bf4e..99e1ed5 100644
--- a/test/unit/test_patterns.py
+++ b/test/unit/test_patterns.py
@@ -25,12 +25,13 @@ from ..script_loader import load_script
 def patterns_code():
     yield load_script('common/patterns.js', ['common'])
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_regexes(execute_in_page, patterns_code):
     """
     patterns.js contains regexes used for URL parsing.
     Verify they work properly.
     """
-    execute_in_page(patterns_code, page='https://gotmyowndoma.in')
+    execute_in_page(patterns_code)
 
     valid_url = 'https://example.com/a/b?ver=1.2.3#heading2'
     valid_url_rest = 'example.com/a/b?ver=1.2.3#heading2'
@@ -90,12 +91,13 @@ def test_regexes(execute_in_page, patterns_code):
                             '@bad.url/')
     assert match is None
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_deconstruct_url(execute_in_page, patterns_code):
     """
     patterns.js contains deconstruct_url() function that handles URL parsing.
     Verify it works properly.
     """
-    execute_in_page(patterns_code, page='https://gotmyowndoma.in')
+    execute_in_page(patterns_code)
 
     deco = execute_in_page('returnval(deconstruct_url(arguments[0]));',
                            'https://eXaMpLe.com/a/b?ver=1.2.3#heading2')
diff --git a/test/unit/test_patterns_query_tree.py b/test/unit/test_patterns_query_tree.py
index e282592..a67e22f 100644
--- a/test/unit/test_patterns_query_tree.py
+++ b/test/unit/test_patterns_query_tree.py
@@ -25,13 +25,14 @@ from ..script_loader import load_script
 def patterns_tree_code():
     yield load_script('common/patterns_query_tree.js', ['common'])
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_modify_branch(execute_in_page, patterns_tree_code):
     """
     patterns_query_tree.js contains Pattern Tree data structure that allows
     arrays of string labels to be mapped to items.
     Verify operations modifying a single branch of such tree work properly.
     """
-    execute_in_page(patterns_tree_code, page='https://gotmyowndoma.in')
+    execute_in_page(patterns_tree_code)
     execute_in_page(
         '''
         let items_added;
@@ -195,13 +196,14 @@ def test_modify_branch(execute_in_page, patterns_tree_code):
                 }
             }
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_search_branch(execute_in_page, patterns_tree_code):
     """
     patterns_query_tree.js contains Pattern Tree data structure that allows
     arrays of string labels to be mapped to items.
     Verify searching a single branch of such tree work properly.
     """
-    execute_in_page(patterns_tree_code, page='https://gotmyowndoma.in')
+    execute_in_page(patterns_tree_code)
     execute_in_page(
         '''
         const item_adder = item => (array => [...(array || []), item]);
@@ -282,13 +284,14 @@ def test_search_branch(execute_in_page, patterns_tree_code):
                   '\nresult:', result, file=sys.stderr)
             raise e from None
 
+@pytest.mark.get_page('https://gotmyowndoma.in')
 def test_pattern_tree(execute_in_page, patterns_tree_code):
     """
     patterns_query_tree.js contains Pattern Tree data structure that allows
     arrays of string labels to be mapped to items.
     Verify operations on entire such tree work properly.
     """
-    execute_in_page(patterns_tree_code, page='https://gotmyowndoma.in')
+    execute_in_page(patterns_tree_code)
 
     # Perform tests with all possible patterns for a simple URL.
     url = 'https://example.com'