1 files changed, 13 insertions, 8 deletions
diff --git a/test/script_loader.py b/test/script_loader.py
index 22196c3..15269c7 100644
--- a/test/script_loader.py
+++ b/test/script_loader.py
@@ -49,14 +49,15 @@ def available_scripts(directory):
         if script_name_regex.match(script.name):
             yield script
 
-def get_wrapped_script(script_path):
+def wrapped_script(script_path, wrap_partially=True):
     if script_path == 'exports_init.js':
         with open(script_root / 'MOZILLA_exports_init.js') as script:
             return script.read()
 
-    awk = subprocess.run(['awk', '-f', str(awk_script), 'wrapped_code',
-                          str(script_path)],
-                         stdout=subprocess.PIPE, cwd=script_root, check=True)
+    command = 'partially_wrapped_code' if wrap_partially else 'wrapped_code'
+    awk_command = ['awk', '-f', str(awk_script), command, str(script_path)]
+    awk = subprocess.run(awk_command, stdout=subprocess.PIPE, cwd=script_root,
+                         check=True)
 
     return awk.stdout.decode()
 
@@ -67,8 +68,10 @@ def load_script(path, import_dirs):
     project directory.
 
     Return a string containing script from `path` together with all other
-    scripts it depends on, wrapped in the same way Haketilo's build system wraps
-    them, with imports properly satisfied.
+    scripts it depends. Dependencies are wrapped in the same way Haketilo's
+    build system wraps them, with imports properly satisfied. The main script
+    being loaded is wrapped partially - it also has its imports satisfied, but
+    its code is not placed inside an anonymous function, so the
     """
     path = make_relative_path(path)
 
@@ -79,6 +82,8 @@ def load_script(path, import_dirs):
                           str(path), *[str(s) for s in available]],
                          stdout=subprocess.PIPE, cwd=script_root, check=True)
 
-    output = awk.stdout.decode()
+    to_load = awk.stdout.decode().split()
+    texts = [wrapped_script(path, wrap_partially=(i == len(to_load) - 1))
+             for i, path in enumerate(to_load)]
 
-    return '\n'.join([get_wrapped_script(path) for path in output.split()])
+    return '\n'.join(texts)