Fix and refactor tests of Docker executors (#827)

tests/test_remote_executors.py

@@ -1,10 +1,11 @@
-import logging
-from unittest import TestCase
+from textwrap import dedent
 from unittest.mock import MagicMock, patch
 
 import docker
+import pytest
 from PIL import Image
 
+from smolagents.monitoring import AgentLogger, LogLevel
 from smolagents.remote_executors import DockerExecutor, E2BExecutor
 
 from .utils.markers import require_run_all
@@ -23,19 +24,23 @@ def test_e2b_executor_instantiation(self):
         assert executor.sandbox == mock_sandbox.return_value
 
 
-class TestDockerExecutor(TestCase):
-    def setUp(self):
-        self.logger = logging.getLogger("DockerExecutorTest")
-        self.executor = DockerExecutor(
-            additional_imports=["pillow", "numpy"], tools=[], logger=self.logger, initial_state={}
-        )
+@pytest.fixture
+def docker_executor():
+    executor = DockerExecutor(additional_imports=["pillow", "numpy"], logger=AgentLogger(level=LogLevel.OFF))
+    yield executor
+    executor.delete()
+
+
+@require_run_all
+class TestDockerExecutor:
+    @pytest.fixture(autouse=True)
+    def set_executor(self, docker_executor):
+        self.executor = docker_executor
 
-    @require_run_all
     def test_initialization(self):
         """Check if DockerExecutor initializes without errors"""
-        self.assertIsNotNone(self.executor.container, "Container should be initialized")
+        assert self.executor.container is not None, "Container should be initialized"
 
-    @require_run_all
     def test_state_persistence(self):
         """Test that variables and imports form one snippet persist in the next"""
         code_action = "import numpy as np; a = 2"
@@ -45,36 +50,42 @@ def test_state_persistence(self):
         result, logs, final_answer = self.executor(code_action)
         assert "1.41421" in logs
 
-    @require_run_all
-    def test_execute_image_output(self):
-        """Test execution that returns a base64 image"""
-        code_action = """
-import base64
-from PIL import Image
-from io import BytesIO
+    def test_execute_output(self):
+        """Test execution that returns a string"""
+        code_action = 'final_answer("This is the final answer")'
+        result, logs, final_answer = self.executor(code_action)
+        assert result == "This is the final answer", "Result should be 'This is the final answer'"
 
-image = Image.new("RGB", (10, 10), (255, 0, 0))
-final_answer(image)
-"""
+    def test_execute_multiline_output(self):
+        """Test execution that returns a string"""
+        code_action = 'result = "This is the final answer"\nfinal_answer(result)'
         result, logs, final_answer = self.executor(code_action)
+        assert result == "This is the final answer", "Result should be 'This is the final answer'"
 
-        self.assertIsInstance(result, Image.Image, "Result should be a PIL Image")
+    def test_execute_image_output(self):
+        """Test execution that returns a base64 image"""
+        code_action = dedent("""
+            import base64
+            from PIL import Image
+            from io import BytesIO
+            image = Image.new("RGB", (10, 10), (255, 0, 0))
+            final_answer(image)
+        """)
+        result, logs, final_answer = self.executor(code_action)
+        assert isinstance(result, Image.Image), "Result should be a PIL Image"
 
-    @require_run_all
     def test_syntax_error_handling(self):
         """Test handling of syntax errors"""
         code_action = 'print("Missing Parenthesis'  # Syntax error
-        with self.assertRaises(ValueError) as context:
+        with pytest.raises(RuntimeError) as exception_info:
             self.executor(code_action)
+        assert "SyntaxError" in str(exception_info.value), "Should raise a syntax error"
 
-        self.assertIn("SyntaxError", str(context.exception), "Should raise a syntax error")
-
-    @require_run_all
     def test_cleanup_on_deletion(self):
         """Test if Docker container stops and removes on deletion"""
         container_id = self.executor.container.id
         self.executor.delete()  # Trigger cleanup
 
         client = docker.from_env()
         containers = [c.id for c in client.containers.list(all=True)]
-        self.assertNotIn(container_id, containers, "Container should be removed")
+        assert container_id not in containers, "Container should be removed"