diff --git a/tests/scripts/generate_code.pl b/tests/scripts/generate_code.pl
index 318629f..6b195e7 100755
--- a/tests/scripts/generate_code.pl
+++ b/tests/scripts/generate_code.pl
@@ -6,11 +6,11 @@
 my $suite_dir = shift or die "Missing suite directory";
 my $suite_name = shift or die "Missing suite name";
 my $data_name = shift or die "Missing data name";
+my $test_main_file = do { my $arg = shift; defined($arg) ? $arg :  $suite_dir."/main_test.function" };
 my $test_file = $data_name.".c";
 my $test_helper_file = $suite_dir."/helpers.function";
 my $test_case_file = $suite_dir."/".$suite_name.".function";
 my $test_case_data = $suite_dir."/".$data_name.".data";
-my $test_main_file = $suite_dir."/main_test.function";
 
 my $line_separator = $/;
 undef $/;