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:(.*)$/;
diff --git a/tests/suites/main_test.function b/tests/suites/main_test.function
index 95a924f..318ca9a 100644
--- a/tests/suites/main_test.function
+++ b/tests/suites/main_test.function
@@ -29,7 +29,7 @@
 
 #define TEST_ASSERT( TEST )                         \
         do { test_assert( (TEST) ? 1 : 0, #TEST );  \
-             if( test_errors) return;               \
+             if( test_errors) goto exit;            \
         } while (0)
 
 int verify_string( char **str )