diff --git a/test/coding/test_embedded_ipython_code_executor.py b/test/coding/test_embedded_ipython_code_executor.py index 75d827fdf518..529925d91bbf 100644 --- a/test/coding/test_embedded_ipython_code_executor.py +++ b/test/coding/test_embedded_ipython_code_executor.py @@ -2,7 +2,7 @@ import sys from pathlib import Path import tempfile -from typing import Dict, Union +from typing import Dict, Type, Union import uuid import pytest from autogen.agentchat.conversable_agent import ConversableAgent @@ -175,6 +175,23 @@ def test_save_image(cls) -> None: assert f"Image data saved to {code_result.output_files[0]}" in code_result.output +@pytest.mark.skipif(skip, reason=skip_reason) +@pytest.mark.parametrize("cls", classes_to_test) +def test_timeout_preserves_kernel_state(cls: Type[CodeExecutor]) -> None: + executor = cls(timeout=1) + code_blocks = [CodeBlock(code="x = 123", language="python")] + code_result = executor.execute_code_blocks(code_blocks) + assert code_result.exit_code == 0 and code_result.output.strip() == "" + + code_blocks = [CodeBlock(code="import time; time.sleep(10)", language="python")] + code_result = executor.execute_code_blocks(code_blocks) + assert code_result.exit_code != 0 and "Timeout" in code_result.output + + code_blocks = [CodeBlock(code="print(x)", language="python")] + code_result = executor.execute_code_blocks(code_blocks) + assert code_result.exit_code == 0 and "123" in code_result.output + + @pytest.mark.skipif(skip, reason=skip_reason) @pytest.mark.parametrize("cls", classes_to_test) def test_save_html(cls) -> None: