diff --git a/tests/scripts/generate_code.pl b/tests/scripts/generate_code.pl
index 5440ced..318629f 100755
--- a/tests/scripts/generate_code.pl
+++ b/tests/scripts/generate_code.pl
@@ -86,7 +86,7 @@
     {
         die "Test function does not have 'void' as return type\n";
     }
-    if ($function_decl !~ /^void (\w+)\(\s*(.*?)\s*\)\s*{(.*?)}/ms)
+    if ($function_decl !~ /^void (\w+)\(\s*(.*?)\s*\)\s*{(.*)}/ms)
     {
         die "Function declaration not in expected format\n";
     }
@@ -104,6 +104,12 @@
 
     $function_decl =~ s/^void /void test_suite_/;
 
+    # Add exit label if not present
+    if ($function_decl !~ /^exit:$/m)
+    {
+        $function_decl =~ s/}\s*$/\nexit:\n    return;\n}/;
+    }
+
     if ($function_deps =~ /^depends_on:/)
     {
         ( $function_deps ) = $function_deps =~ /^depends_on:(.*)$/;