Keep output message when program exit in PDBTool.

Copy from #122, thanks Marc :)

Author: matheper

debug_gym/gym/tools/pdb.py

@@ -181,12 +181,11 @@ def use(self, environment, command: str) -> Observation:
         # into self.last_eval_output, and remove them from the output
         if "The program exited via sys.exit()." in output:
             # end index is the last occurrence of the program exited (from the \n after)
-            end_index = (
-                output.find("\n", output.rfind("The program exited via sys.exit()."))
-                + 1
-            )
+            start_index = output.rfind("The program exited via sys.exit().")
+            end_index = output.find("\n", start_index) + 1
             output = (
-                "Reached the end of the file. Restarting the debugging session.\n"
+                output[:start_index]
+                + "\nReached the end of the program. Restarting the debugging session.\n"
                 + output[end_index:]
             )
         obs = "\n".join([_warning, output]).strip() + "\n"

tests/gym/tools/test_pdb.py

@@ -77,10 +77,13 @@ def test_pdb_use(tmp_path, setup_test_repo):
     assert """The pytest entry point.""" in output
     assert "(Pdb)" not in output
     output = pdb.use(environment, command="c").observation
-    assert "1 failed, 1 passed" in pdb.pdb_obs
-    assert "test_fail.py::test_fail FAILED" in pdb.pdb_obs
-    assert "test_pass.py::test_pass PASSED" in pdb.pdb_obs
-    assert "Reached the end of the file. Restarting the debugging session." in output
+    assert "1 failed, 1 passed" in output
+    assert "test_fail.py::test_fail FAILED" in output
+    assert "test_pass.py::test_pass PASSED" in output
+    assert "Reached the end of the program. Restarting the debugging session." in output
+    assert "pytest/__main__.py" in output
+    assert '-> """The pytest entry point."""' in output
+    assert 'Context around the current frame:\n1  ->	"""The pytest entry point.""""'
     assert "(Pdb)" not in output
 
 
@@ -152,10 +155,13 @@ def test_pdb_use_default_environment_entrypoint(tmp_path, setup_test_repo):
     assert "(Pdb)" not in output
 
     output = pdb.use(environment, command="c").observation
-    assert "1 failed, 1 passed" in pdb.pdb_obs
-    assert "test_fail.py::test_fail" in pdb.pdb_obs
-    assert "test_pass.py::test_pass" not in pdb.pdb_obs
-    assert "Reached the end of the file. Restarting the debugging session." in output
+    assert "1 failed, 1 passed" in output
+    assert "test_fail.py::test_fail" in output
+    assert "test_pass.py::test_pass" not in output
+    assert "Reached the end of the program. Restarting the debugging session." in output
+    assert "pytest/__main__.py" in output
+    assert '-> """The pytest entry point."""' in output
+    assert 'Context around the current frame:\n1  ->	"""The pytest entry point.""""'
     assert "(Pdb)" not in output
 
 
@@ -182,10 +188,13 @@ def test_pdb_use_docker_terminal(tmp_path, setup_test_repo):
     assert "(Pdb)" not in output
 
     output = pdb.use(environment, command="c").observation
-    assert "1 failed, 1 passed" in pdb.pdb_obs
-    assert "test_fail.py::test_fail FAILED" in pdb.pdb_obs
-    assert "test_pass.py::test_pass PASSED" in pdb.pdb_obs
-    assert "Reached the end of the file. Restarting the debugging session." in output
+    assert "1 failed, 1 passed" in output
+    assert "test_fail.py::test_fail FAILED" in output
+    assert "test_pass.py::test_pass PASSED" in output
+    assert "Reached the end of the program. Restarting the debugging session." in output
+    assert "pytest/__main__.py" in output
+    assert '-> """The pytest entry point."""' in output
+    assert 'Context around the current frame:\n1  ->	"""The pytest entry point.""""'
     assert "(Pdb)" not in output