diff --git a/yaksh/evaluator_tests/test_java_evaluation.py b/yaksh/evaluator_tests/test_java_evaluation.py
index eb09f2f15..5779cc5c6 100644
--- a/yaksh/evaluator_tests/test_java_evaluation.py
+++ b/yaksh/evaluator_tests/test_java_evaluation.py
@@ -20,43 +20,8 @@ def setUp(self):
             f.write('2'.encode('ascii'))
         tmp_in_dir_path = tempfile.mkdtemp()
         self.tc_data = dedent("""
-            class main
-            {
-                public static <E> void check(E expect, E result)
-                {
-                    if(result.equals(expect))
-                    {
-                        System.out.println("Correct:Output expected "+expect+
-                                            "and got "+result);
-                    }
-                else
-                {
-                    System.out.println("Incorrect:Output expected "+expect+
-                                        "but got "+result);
-                    System.exit(1);
-                }
-                }
-                public static void main(String arg[])
-                {
-                   Test t = new Test();
-                   int result, input, output;
-                   input = 0; output = 0;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                    input);
-                   check(output, result);
-                   input = 5; output = 25;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                    input);
-                   check(output, result);
-                   input = 6; output = 36;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                    input);
-                   check(output, result);
-                }
-            }
+            Test t = new Test();
+            assert t.square_num(5) == 25: "Square of 5 is not 25";
             """)
 
         self.test_case_data = [
@@ -172,33 +137,12 @@ def test_file_based_assert(self):
         # Given
         self.file_paths = [(self.f_path, False)]
         self.tc_data = dedent("""
-            class main
-            {
-                public static <E> void check(E expect, E result)
-                {
-                    if(result.equals(expect))
-                    {
-                        System.out.println("Correct:Output expected "+expect+
-                                            " and got "+result);
-                    }
-                else
-                {
-                    System.out.println("Incorrect:Output expected "+expect+
-                                        " but got "+result);
-                    System.exit(1);
-                }
-                }
-                public static void main(String arg[])
-                {
-                   String result = "";
-                   Test t = new Test();
-                   try{
-                   result = t.readFile();}
-                   catch(Exception e){
+            Test t = new Test();
+            try{
+                String a = t.readFile();
+                assert a.equals("2");
+            } catch(Exception e) {
                 System.out.print(e);
-                }
-                   check("2", result);
-                }
             }
             """)
         self.test_case_data = [
@@ -237,40 +181,14 @@ class Test{
         # When
         grader = Grader(self.in_dir)
         result = grader.evaluate(kwargs)
-
         # Then
         self.assertTrue(result.get("success"))
 
     def test_incorrect_testcase(self):
         # Given
         self.tc_data = dedent("""
-            class main
-            {
-                public static <E> void check(E expect, E result)
-                {
-                    if(result.equals(expect))
-                    {
-                        System.out.println("Correct:Output expected "+expect+
-                                            "and got "+result);
-                    }
-                else
-                {
-                    System.out.println("Incorrect:Output expected "+expect+
-                                        "but got "+result);
-                    System.exit(1);
-                }
-                }
-                public static void main(String arg[])
-                {
-                   Test t = new Test();
-                   int result, input, output;
-                   input = 0; output = 0;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                    input);
-                   check(output, result)
-                }
-            }
+            Test t = new Test()
+            result = t.square_num(5) == 25;
             """)
         user_answer = ("class Test {\n\tint square_num(int a) "
                        "{\n\treturn a;\n\t}\n}")
@@ -701,43 +619,8 @@ def test_assert_with_hook(self):
         user_answer = ("class Test {\n\tint square_num(int a)"
                        " {\n\treturn a*a;\n\t}\n}")
         assert_test_case = dedent("""
-            class main
-            {
-                public static <E> void check(E expect, E result)
-                {
-                    if(result.equals(expect))
-                    {
-                        System.out.println("Correct:Output expected "+expect+
-                                            " and got "+result);
-                    }
-                else
-                {
-                    System.out.println("Incorrect:Output expected "+expect+
-                                        " but got "+result);
-                    System.exit(1);
-                }
-                }
-                public static void main(String arg[])
-                {
-                   Test t = new Test();
-                   int result, input, output;
-                   input = 0; output = 0;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                        input);
-                   check(output, result);
-                   input = 5; output = 25;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                        input);
-                   check(output, result);
-                   input = 6; output = 36;
-                   result = t.square_num(input);
-                   System.out.println("Input submitted to the function: "+
-                                        input);
-                   check(output, result);
-                }
-            }
+            Test t = new Test();
+            assert t.square_num(6) == 36;
             """)
 
         hook_code = dedent("""\
diff --git a/yaksh/fixtures/demo_questions.zip b/yaksh/fixtures/demo_questions.zip
index 6b0f8526c..b35499134 100644
Binary files a/yaksh/fixtures/demo_questions.zip and b/yaksh/fixtures/demo_questions.zip differ
diff --git a/yaksh/java_code_evaluator.py b/yaksh/java_code_evaluator.py
index 35573c06d..11dab717b 100644
--- a/yaksh/java_code_evaluator.py
+++ b/yaksh/java_code_evaluator.py
@@ -3,6 +3,7 @@
 import os
 from os.path import isfile
 import subprocess
+from textwrap import dedent
 
 # Local imports
 from .base_evaluator import BaseEvaluator
@@ -53,17 +54,28 @@ def set_file_paths(self, directory, file_name):
         output_path = "{0}{1}.class".format(directory, file_name)
         return output_path
 
+    def create_main_file_content(self):
+        content = dedent('''
+        class main {{
+            public static void main( String args[] ) {{
+               {0};
+            }}
+        }}
+        '''.format(self.test_case))
+        return content
+
     def compile_code(self):
         if self.compiled_user_answer and self.compiled_test_code:
             return None
         else:
             # create student code and moderator code file
-            self.submit_code_path = self.create_submit_code_file('Test.java')
+            self.submit_code_path = self.create_submit_code_file('code.java')
             self.test_code_path = self.create_submit_code_file('main.java')
             self.write_to_submit_code_file(
                 self.submit_code_path, self.user_answer
                 )
-            self.write_to_submit_code_file(self.test_code_path, self.test_case)
+            main_content = self.create_main_file_content()
+            self.write_to_submit_code_file(self.test_code_path, main_content)
             clean_ref_code_path = self.test_code_path
             if self.file_paths:
                 self.files = copy_files(self.file_paths)
@@ -86,7 +98,7 @@ def compile_code(self):
                 clean_ref_code_path,
                 user_code_directory
             )
-            self.run_command_args = "java -cp {0} {1}".format(
+            self.run_command_args = "java -ea -cp {0} {1}".format(
                 user_code_directory,
                 ref_file_name
             )