Fixes an issue where exception text was not being included in command output

Author: pranavm-nvidia

tripy/tests/helper.py

@@ -498,19 +498,18 @@ def process_code_block_for_outputs_and_locals(
     code_start, code_end = get_code_bounds(block_lines)
     code = dedent("\n".join(block_lines[code_start:code_end]))
 
-    try:
-        with capture_output() as outfile:
+    with capture_output() as outfile:
+        try:
             code_locals = exec_code(code, local_vars)
-    except Exception as e:
-        if allow_exception:
-            code_locals = local_vars
-        else:
-            print(
-                f"Exception occurred while executing code block: {type(e).__name__}: {e}\n"
-                f"Note: Code block was:\n\n{block}"
-            )
-            print(err_msg)
-            raise
+        except Exception as e:
+            if allow_exception:
+                # We print the error message here so it can be captured in `outfile`
+                # and displayed in the output in cases where we actually expect exceptions.
+                print(e)
+                code_locals = local_vars
+            else:
+                print(f"{err_msg}\n" f"Note: Code block was:\n\n{block}")
+                raise
 
     new_locals = {
         key: value for key, value in code_locals.items() if key not in local_vars or value is not local_vars[key]