Update nodes.py

nodes.py

@@ -136,26 +136,28 @@ def test_generator(state: State):
 
     response = llm.invoke([
         SystemMessage(content="""You are a Python testing expert.
-Return ONLY runnable Python test code — no markdown, no backticks."""),
+Return ONLY runnable Python test code — no markdown, no backticks.
+DO NOT use 'unittest', 'pytest', or 'sys'."""),
         HumanMessage(content=f"""
 Generate test cases for this code:
-
 TASK: {state['task']}
 CODE:
 {code}
 
 Rules:
-- Copy ALL function definitions inline — do NOT import from files
-- Cover: normal cases, edge cases, large input
-- Call each test function at the bottom to run them
-- Do NOT use unittest or sys — just plain assert statements
-- Print "All tests passed!" at the end if successful
+- Copy ALL function definitions inline.
+- Use ONLY simple 'assert' statements for validation.
+- Do NOT use 'unittest' or 'sys'.
+- If a test fails, let the script raise an AssertionError.
+- Print "All tests passed!" at the end if successful.
+- Wrap all test calls in a 'try...except' block to print the error before exiting.
 
 Return ONLY runnable Python code:
 """)
     ])
 
     tests = response.content
+    # ... (keep existing regex cleaning)
     tests = re.sub(r"```python", "", tests)
     tests = re.sub(r"```", "", tests)
     tests = tests.strip()