1 files changed, 83 insertions, 12 deletions

diff --git a/test/unit/test_content.py b/test/unit/test_content.py
index c8e0987..35ab027 100644
--- a/test/unit/test_content.py
+++ b/test/unit/test_content.py
@@ -42,7 +42,7 @@ dynamic_script = \
 content_script = \
     '''
     /* Mock dynamic content script - case 'before'. */
-    if (/#dynamic_before$/.test(document.URL)) {
+    if (/dynamic_before/.test(document.URL)) {
         %s;
     }
 
@@ -50,6 +50,37 @@ content_script = \
     %s;
 
     /* Rest of mocks */
+
+    function mock_decide_policy() {
+        nonce = "12345";
+        return {
+            allow: false,
+            mapping: "what-is-programmers-favorite-drinking-place",
+            payload: {identifier: "foo-bar"},
+            nonce,
+            csp: "prefetch-src 'none'; script-src-attr 'none'; script-src 'nonce-12345'; script-src-elem 'nonce-12345';"
+        };
+    }
+
+    async function mock_payload_error([type, res_id]) {
+        if (type === "indexeddb_files")
+            return {error: {haketilo_error_type: "missing", id: res_id}};
+    }
+
+    async function mock_payload_ok([type, res_id]) {
+        if (type === "indexeddb_files")
+            return [1, 2].map(n => `window.haketilo_injected_${n} = ${n}${n};`);
+    }
+
+    if (/payload_error/.test(document.URL)) {
+        browser.runtime.sendMessage = mock_payload_error;
+        decide_policy = mock_decide_policy;
+    } else if (/payload_ok/.test(document.URL)) {
+        browser.runtime.sendMessage = mock_payload_ok;
+        decide_policy = mock_decide_policy;
+    }
+    /* Otherwise, script blocking policy without payload to inject is used. */
+
     const data_to_verify = {};
     function data_set(prop, val) {
         data_to_verify[prop] = val;
@@ -61,22 +92,24 @@ content_script = \
     enforce_blocking = policy => data_set("enforcing", policy);
 
     browser.runtime.onMessage.addListener = async function (listener_cb) {
-        await new Promise(cb => setTimeout(cb, 0));
+        await new Promise(cb => setTimeout(cb, 10));
 
         /* Mock a good request. */
         const set_good = val => data_set("good_request_result", val);
-        listener_cb(["page_info"], {}, val => set_good(val));
+        data_set("good_request_returned",
+                 !!listener_cb(["page_info"], {}, val => set_good(val)));
 
         /* Mock a bad request. */
         const set_bad = val => data_set("bad_request_result", val);
-        listener_cb(["???"], {}, val => set_bad(val));
+        data_set("bad_request_returned",
+                 !!listener_cb(["???"], {}, val => set_bad(val)));
     }
 
     /* main() call - normally present in content.js, inside '#IF !UNIT_TEST'. */
     main();
 
     /* Mock dynamic content script - case 'after'. */
-    if (/#dynamic_after$/.test(document.URL)) {
+    if (/#dynamic_after/.test(document.URL)) {
         %s;
     }
 
@@ -85,26 +118,64 @@ content_script = \
 
 @pytest.mark.ext_data({'content_script': content_script})
 @pytest.mark.usefixtures('webextension')
-@pytest.mark.parametrize('target', ['dynamic_before', 'dynamic_after'])
-def test_content_unprivileged_page(driver, execute_in_page, target):
+@pytest.mark.parametrize('target1', ['dynamic_before'])#, 'dynamic_after'])
+@pytest.mark.parametrize('target2', [
+    'scripts_blocked',
+    'payload_error',
+    'payload_ok'
+])
+def test_content_unprivileged_page(driver, execute_in_page, target1, target2):
     """
     Test functioning of content.js on an page using unprivileged schema (e.g.
     'https://' and not 'about:').
     """
-    driver.get(f'https://gotmyowndoma.in/index.html#{target}')
-    data = json.loads(driver.execute_script('return window.data_to_verify;'))
+    driver.get(f'https://gotmyowndoma.in/index.html#{target1}-{target2}')
+
+    def get_data(driver):
+        data = driver.execute_script('return window.data_to_verify;')
+        return data if 'good_request_result' in data else False
+
+    data = json.loads(WebDriverWait(driver, 10).until(get_data))
 
     assert 'gotmyowndoma.in' in data['good_request_result']['url']
     assert 'bad_request_result' not in data
 
+    assert data['good_request_returned'] == True
+    assert data['bad_request_returned'] == False
+
     assert data['cacher_started'] == True
 
-    assert data['enforcing']['allow']  == False
-    assert 'mapping' not in data['enforcing']
-    assert 'error'   not in data['enforcing']
+    for obj in (data['good_request_result'], data['enforcing']):
+        assert obj['allow'] == False
+
+    assert 'error' not in data['enforcing']
+
+    if target2.startswith('payload'):
+        for obj in (data['good_request_result'], data['enforcing']):
+            assert obj['payload']['identifier'] == 'foo-bar'
+            assert 'mapping' in obj
+    else:
+        assert 'payload' not in data['enforcing']
+        assert 'mapping' not in data['enforcing']
 
     assert data['script_run_without_errors'] == True
 
+    def vars_made_by_payload(driver):
+        vars_values = driver.execute_script(
+            'return [1, 2].map(n => window[`haketilo_injected_${n}`]);'
+        )
+        if vars_values != [None, None]:
+            return vars_values
+
+    if target2 == 'payload_error':
+        assert data['good_request_result']['error'] == {
+            'haketilo_error_type': 'missing',
+            'id': 'foo-bar'
+        }
+    elif target2 == 'payload_ok':
+        vars_values = WebDriverWait(driver, 10).until(vars_made_by_payload)
+        assert vars_values == [11, 22]
+
 @pytest.mark.ext_data({'content_script': content_script})
 @pytest.mark.usefixtures('webextension')
 @pytest.mark.parametrize('target', ['dynamic_before', 'dynamic_after'])