1 files changed, 31 insertions, 13 deletions

diff --git a/test/unit/test_patterns_query_manager.py b/test/unit/test_patterns_query_manager.py
index 4e6d1bf..c6ebb81 100644
--- a/test/unit/test_patterns_query_manager.py
+++ b/test/unit/test_patterns_query_manager.py
@@ -18,7 +18,6 @@ Haketilo unit tests - building pattern tree and putting it in a content script
 # CC0 1.0 Universal License for more details.
 
 import pytest
-import re
 import json
 from selenium.webdriver.support.ui import WebDriverWait
 
@@ -35,13 +34,19 @@ def simple_sample_mapping(patterns, fruit):
         'payloads': payloads
     }
 
-content_script_tree_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
-def extract_tree_data(content_script_text):
-    return json.loads(content_script_tree_re.search(content_script_text)[1])
-
-content_script_mapping_re = re.compile(r'this.haketilo_mappings = (.*);')
-def extract_mappings_data(content_script_text):
-    return json.loads(content_script_mapping_re.search(content_script_text)[1])
+def get_content_script_values(driver, content_script):
+    """
+    Allow easy extraction of 'this.something = ...' values from generated
+    content script and verify the content script is syntactically correct.
+    """
+    return driver.execute_script(
+        '''
+        function value_holder() {
+            %s;
+            return this;
+        }
+        return value_holder.call({});
+        ''' % content_script)
 
 # Fields that are not relevant for testing are omitted from these mapping
 # definitions.
@@ -87,7 +92,7 @@ def test_pqm_tree_building(driver, execute_in_page):
     execute_in_page(
         '''
         const [initial_mappings, initial_blocking] = arguments.slice(0, 2);
-        let mappingchange, blockingchange;
+        let mappingchange, blockingchange, settingchange;
 
         haketilodb.track.mapping = function (cb) {
             mappingchange = cb;
@@ -99,6 +104,11 @@ def test_pqm_tree_building(driver, execute_in_page):
 
             return [{}, initial_blocking];
         }
+        haketilodb.track.settings = function (cb) {
+            settingchange = cb;
+
+            return [{}, [{name: "default_allow", value: true}]];
+        }
 
         let last_script;
         let unregister_called = 0;
@@ -110,7 +120,7 @@ def test_pqm_tree_building(driver, execute_in_page):
         }
         browser = {contentScripts: {register: register_mock}};
 
-        returnval(start());
+        returnval(start("abracadabra"));
         ''',
         sample_mappings[0:2], sample_blocking[0:2])
 
@@ -125,17 +135,24 @@ def test_pqm_tree_building(driver, execute_in_page):
         dict([('~allow', 1),
               *[(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
                 for fruit in ('banana', 'orange')]])
-    assert tree == extract_tree_data(content_script)
+    cs_values = get_content_script_values(driver, content_script)
+    assert cs_values['haketilo_secret']        == 'abracadabra'
+    assert cs_values['haketilo_pattern_tree']  == tree
+    assert cs_values['haketilo_default_allow'] == True
     assert deregistrations == 0
 
     def condition_all_added(driver):
         last_script = execute_in_page('returnval(last_script);')
+        cs_values = get_content_script_values(driver, last_script)
         nums = [i for i in range(len(sample_blocking)) if i > 1]
-        return (all([('gotmyown%sdoma' % i) in last_script for i in nums]) and
+        return (cs_values['haketilo_default_allow'] == False and
+                all([('gotmyown%sdoma' % i) in last_script for i in nums]) and
                 all([m['identifier'] in last_script for m in sample_mappings]))
 
     execute_in_page(
         '''
+        const new_setting_val = {name: "default_allow", value: false};
+        settingchange({key: "default_allow", new_val: new_setting_val});
         for (const mapping of arguments[0])
             mappingchange({key: mapping.identifier, new_val: mapping});
         for (const blocking of arguments[1])
@@ -163,7 +180,8 @@ def test_pqm_tree_building(driver, execute_in_page):
 
     def condition_all_removed(driver):
         content_script = execute_in_page('returnval(last_script);')
-        return extract_tree_data(content_script) == {}
+        cs_values = get_content_script_values(driver, content_script)
+        return cs_values['haketilo_pattern_tree'] == {}
 
     execute_in_page(
         '''