diff --git a/src/itp_interface/tools/lean4_sync_executor.py b/src/itp_interface/tools/lean4_sync_executor.py
index f3e7ae2..1e7ba6c 100644
--- a/src/itp_interface/tools/lean4_sync_executor.py
+++ b/src/itp_interface/tools/lean4_sync_executor.py
@@ -771,7 +771,7 @@ def _parse_response(self, idx, response, relevant_messages = [], dont_update_err
             messages = response['messages']
             # Go over all sev after the line number and check if there is an error
             for msg_idx, msg in enumerate(messages):
-                full_error_msg = self._get_error_msg(msg_idx, msg)
+                full_error_msg = f"Line: {idx} " + self._get_error_msg(msg_idx, msg)
                 unsolved_goal_never_seen_before = not (full_error_msg in self._error_messages_since_last_thm.values())
                 if msg['severity'] == 'error' and 'pos' in msg and 'endPos' in msg and \
                 ((msg['endPos'] is not None and 'line' in msg['endPos']) or \
diff --git a/src/test/simple_env_test.py b/src/test/simple_env_test.py
index d6f3b1d..62b9fea 100644
--- a/src/test/simple_env_test.py
+++ b/src/test/simple_env_test.py
@@ -80,13 +80,17 @@ def test_simple_lean4(self):
         retrieval_strategy = ProofEnvReRankStrategy.NO_RE_RANK
         env = ProofEnv("test_lean4", proof_exec_callback, theorem_name, retrieval_strategy=retrieval_strategy, max_proof_depth=10, always_retrieve_thms=always_retrieve_thms)
         proof_steps = [
+            '-- TODO',
             'apply And.intro',
             'exact hp',
             'apply And.intro',
+            '--TODO',
+            '-- This is just some extra comment',
             'exact hq',
             'exact hp'
         ]
         with env:
+            proof_was_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -99,6 +103,7 @@ def test_simple_lean4(self):
                 print('-'*30)
                 if done:
                     print("Proof Finished!!")
+                    proof_was_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -120,6 +125,7 @@ def test_simple_lean4(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_was_finished, "Proof was not finished"
 
     def test_lean4_backtracking(self):
         from itp_interface.rl.proof_state import ProofState
@@ -156,6 +162,7 @@ def test_lean4_backtracking(self):
         ]
         with env:
             prev_state = env.state
+            proof_was_finished = False
             for idx, proof_step in enumerate(proof_steps):
                 if idx > 0 and random.random() <= 0.5:
                     print(f"Backtracking at step {idx + 1} i.e. {proof_step}")
@@ -179,6 +186,8 @@ def test_lean4_backtracking(self):
                 prev_state = state
                 if done:
                     print("Proof Finished!!")
+                    proof_was_finished = True
+            assert proof_was_finished, "Proof was not finished"
 
     def test_simple_coq(self):
         from itp_interface.rl.proof_state import ProofState
@@ -286,6 +295,7 @@ def test_simple_lean_calc(self):
 "_ = (n + 1)*(n + 1) := by \n   rw [Nat.right_distrib n 1 (n + 1)]"
 ]
         with env:
+            proof_was_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -310,6 +320,7 @@ def test_simple_lean_calc(self):
                     print(f"Action: {proof_step}")
                     print(f"="*30)
                     print("Proof Finished!!")
+                    proof_was_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -331,6 +342,7 @@ def test_simple_lean_calc(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_was_finished, "Proof was not finished"
 
     def test_simple_lean_enforce_done_test(self):
         from itp_interface.rl.proof_state import ProofState
@@ -372,6 +384,7 @@ def test_simple_lean_enforce_done_test(self):
 "done"
 ]
         with env:
+            proof_finished = False
             for proof_step in proof_steps:
                 state, _, next_state, _, done, info = env.step(ProofAction(
                     ProofAction.ActionType.RUN_TACTIC, 
@@ -397,6 +410,7 @@ def test_simple_lean_enforce_done_test(self):
                     print(f"Action: {proof_step}")
                     print(f"="*30)
                     print("Proof Finished!!")
+                    proof_finished = True
                 else:
                     s1 : ProofState = state
                     s2 : ProofState = next_state
@@ -418,6 +432,7 @@ def test_simple_lean_enforce_done_test(self):
                         print('|- ', end='')
                         print(goal.goal)
                     print(f"="*30)
+            assert proof_finished, "Proof was not finished"
 
     def test_simple_lean4_done_test(self):
         from itp_interface.rl.proof_state import ProofState