1 files changed, 21 insertions, 18 deletions

diff --git a/test/unit/test_patterns_query_manager.py b/test/unit/test_patterns_query_manager.py
index 8ae7c28..ae1f490 100644
--- a/test/unit/test_patterns_query_manager.py
+++ b/test/unit/test_patterns_query_manager.py
@@ -25,10 +25,9 @@ from selenium.webdriver.support.ui import WebDriverWait
 from ..script_loader import load_script
 
 def simple_sample_mapping(patterns, fruit):
-    if type(patterns) is list:
-        payloads = dict([(p, {'identifier': fruit}) for p in patterns])
-    else:
-        payloads = {patterns: {'identifier': fruit}}
+    if type(patterns) is not list:
+        patterns = [patterns]
+    payloads = dict([(p, {'identifier': f'{fruit}-{p}'}) for p in patterns])
     return {
         'source_copyright': [],
         'type': 'mapping',
@@ -36,9 +35,13 @@ def simple_sample_mapping(patterns, fruit):
         'payloads': payloads
     }
 
-content_script_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
+content_script_tree_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
 def extract_tree_data(content_script_text):
-    return json.loads(content_script_re.search(content_script_text)[1])
+    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])
 
 # Fields that are not relevant for testing are omitted from these mapping
 # definitions.
@@ -82,7 +85,7 @@ def test_pqm_tree_building(driver, execute_in_page):
 
             return [{}, initial_mappings];
         }
-        haketilodb.track_mappings = track_mock;
+        haketilodb.track.mappings = track_mock;
 
         let last_script;
         let unregister_called = 0;
@@ -104,7 +107,10 @@ def test_pqm_tree_building(driver, execute_in_page):
                    tree, last_script, unregister_called]);
         ''',
         'https://gotmyowndoma.in/index.html')
-    assert found == dict([(m['identifier'], m) for m in sample_mappings[0:2]])
+    best_pattern = 'https://gotmyowndoma.in/index.html'
+    assert found == \
+        dict([(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
+              for fruit in ('banana', 'orange')])
     assert tree == extract_tree_data(content_script)
     assert deregistrations == 0
 
@@ -114,12 +120,8 @@ def test_pqm_tree_building(driver, execute_in_page):
 
     execute_in_page(
         '''
-        for (const mapping of arguments[0]) {
-            mappingchange({
-                identifier: mapping.identifier,
-                new_val: mapping
-            });
-        }
+        for (const mapping of arguments[0])
+            mappingchange({key: mapping.identifier, new_val: mapping});
         ''',
         sample_mappings[2:])
     WebDriverWait(driver, 10).until(condition_mappings_added)
@@ -129,7 +131,8 @@ def test_pqm_tree_building(driver, execute_in_page):
 
     def condition_odd_removed(driver):
         last_script = execute_in_page('returnval(last_script);')
-        return all([id not in last_script for id in odd])
+        return (all([id not in last_script for id in odd]) and
+                all([id in last_script for id in even]))
 
     def condition_all_removed(driver):
         content_script = execute_in_page('returnval(last_script);')
@@ -137,7 +140,7 @@ def test_pqm_tree_building(driver, execute_in_page):
 
     execute_in_page(
         '''
-        arguments[0].forEach(identifier => mappingchange({identifier}));
+        arguments[0].forEach(identifier => mappingchange({key: identifier}));
         ''',
         odd)
 
@@ -145,7 +148,7 @@ def test_pqm_tree_building(driver, execute_in_page):
 
     execute_in_page(
         '''
-        arguments[0].forEach(identifier => mappingchange({identifier}));
+        arguments[0].forEach(identifier => mappingchange({key: identifier}));
         ''',
         even)
 
@@ -224,7 +227,7 @@ def test_pqm_script_injection(driver, execute_in_page):
             const ctx = await start_items_transaction(["mappings"], {});
             for (const id of identifiers)
                 await remove_mapping(id, ctx);
-            await finalize_items_transaction(ctx);
+            await finalize_transaction(ctx);
         }
         returnval(remove_items());
         }''',