Fix linting and tests

Author: Joffref

src/smolagents/remote_executors.py

@@ -660,16 +660,8 @@ def __init__(
                 "name": self.sandbox_name,
             },
             "spec": {
-                "runtime": {
-                    "image": image,
-                    "memory": memory,
-                    "ports": [
-                        {
-                            "target": self.port
-                        }
-                    ]
-                },
-            }
+                "runtime": {"image": image, "memory": memory, "ports": [{"target": self.port}]},
+            },
         }
 
         if region:
@@ -685,7 +677,12 @@ def __init__(
 
             # Create kernel via HTTP
             from blaxel.core import settings
-            kernel_id = _create_kernel_http(f"{self.sandbox.metadata.url}/port/{self.port}/api/kernels?token={token}", self.logger, headers=settings.headers)
+
+            kernel_id = _create_kernel_http(
+                f"{self.sandbox.metadata.url}/port/{self.port}/api/kernels?token={token}",
+                self.logger,
+                headers=settings.headers,
+            )
 
             # Set up websocket URL
             # Convert http/https to ws/wss
@@ -695,22 +692,21 @@ def __init__(
 
             # Install additional packages
             self.installed_packages = self.install_packages(additional_imports)
-            self.logger.log(f"Blaxel is running", level=LogLevel.INFO)
+            self.logger.log("Blaxel is running", level=LogLevel.INFO)
         except Exception as e:
             self.cleanup()
             raise RuntimeError(f"Failed to initialize Blaxel sandbox: {e}") from e
 
     @staticmethod
     def _create_sandbox(config):
         """Helper method to create sandbox asynchronously."""
+        from blaxel.core import SandboxInstance
         from blaxel.core.client import client
         from blaxel.core.client.api.compute import create_sandbox
-        from blaxel.core import SandboxInstance
 
         response = create_sandbox.sync(client=client, body=config)
         return SandboxInstance(response)
 
-
     def run_code_raise_errors(self, code: str) -> CodeOutput:
         """
         Execute Python code in the Blaxel sandbox and return the result.
@@ -721,8 +717,8 @@ def run_code_raise_errors(self, code: str) -> CodeOutput:
         Returns:
             `CodeOutput`: Code output containing the result, logs, and whether it is the final answer.
         """
-        from websocket import create_connection
         from blaxel.core import settings
+        from websocket import create_connection
 
         headers = []
         for key, value in settings.headers.items():
@@ -735,10 +731,10 @@ def install_packages(self, additional_imports: list[str]) -> list[str]:
         if not additional_imports:
             return []
 
-        from blaxel.core.sandbox.client import client
         from blaxel.core import settings
-        from blaxel.core.sandbox.client.models import ProcessResponse, ErrorResponse
-        from blaxel.core.sandbox.client.api.process import post_process, get_process_identifier
+        from blaxel.core.sandbox.client import client
+        from blaxel.core.sandbox.client.api.process import get_process_identifier, post_process
+        from blaxel.core.sandbox.client.models import ErrorResponse, ProcessResponse
 
         try:
             client.with_base_url(self.sandbox.metadata.url)

tests/test_remote_executors.py

@@ -565,61 +565,112 @@ def test_blaxel_executor_instantiation_without_blaxel_sdk(self):
                 BlaxelExecutor(additional_imports=[], logger=logger)
             assert "Please install 'blaxel' extra" in str(excinfo.value)
 
+    @patch("smolagents.remote_executors._create_kernel_http")
     @patch("blaxel.core.SandboxInstance")
-    def test_blaxel_executor_instantiation_with_blaxel_sdk(self, mock_sandbox_instance):
+    @patch("blaxel.core.client.api.compute.create_sandbox")
+    @patch("blaxel.core.settings")
+    def test_blaxel_executor_instantiation_with_blaxel_sdk(
+        self, mock_settings, mock_create_sandbox, mock_sandbox_instance, mock_create_kernel
+    ):
         """Test BlaxelExecutor instantiation with mocked Blaxel SDK."""
         logger = MagicMock()
+        mock_settings.headers = {}
+        
+        # Mock sandbox response
+        mock_response = MagicMock()
+        mock_create_sandbox.sync.return_value = mock_response
+        
+        # Mock SandboxInstance
         mock_sandbox = MagicMock()
-        mock_sandbox_instance.create.return_value = mock_sandbox
-
-        with patch("asyncio.run") as mock_asyncio_run:
-            mock_asyncio_run.return_value = mock_sandbox
-            executor = BlaxelExecutor(additional_imports=[], logger=logger)
-
-            assert executor.sandbox_name == "smolagent-executor"
-            assert executor.image == "blaxel/prod-py-app:latest"
-            assert executor.memory == 4096
-            assert executor.region is None
-
+        mock_metadata = MagicMock()
+        mock_metadata.url = "https://test-sandbox.bl.run"
+        mock_sandbox.metadata = mock_metadata
+        mock_sandbox_instance.return_value = mock_sandbox
+        
+        # Mock kernel creation
+        mock_create_kernel.return_value = "kernel-123"
+        
+        executor = BlaxelExecutor(additional_imports=[], logger=logger)
+
+        assert executor.sandbox_name.startswith("smolagent-executor-")
+        assert executor.image == "blaxel/jupyter-notebook"
+        assert executor.memory == 4096
+        assert executor.region is None
+
+    @patch("smolagents.remote_executors.BlaxelExecutor.install_packages")
+    @patch("smolagents.remote_executors._create_kernel_http")
     @patch("blaxel.core.SandboxInstance")
-    def test_blaxel_executor_custom_parameters(self, mock_sandbox_instance):
+    @patch("blaxel.core.client.api.compute.create_sandbox")
+    @patch("blaxel.core.settings")
+    def test_blaxel_executor_custom_parameters(
+        self, mock_settings, mock_create_sandbox, mock_sandbox_instance, mock_create_kernel, mock_install_packages
+    ):
         """Test BlaxelExecutor with custom parameters."""
         logger = MagicMock()
+        mock_settings.headers = {}
+        mock_install_packages.return_value = ["numpy"]
+        
+        # Mock sandbox response
+        mock_response = MagicMock()
+        mock_create_sandbox.sync.return_value = mock_response
+        
+        # Mock SandboxInstance
         mock_sandbox = MagicMock()
-        mock_sandbox_instance.create.return_value = mock_sandbox
-
-        with patch("asyncio.run") as mock_asyncio_run:
-            mock_asyncio_run.return_value = mock_sandbox
-            executor = BlaxelExecutor(
-                additional_imports=["numpy"],
-                logger=logger,
-                sandbox_name="test-sandbox",
-                image="custom-image:latest",
-                memory=8192,
-                region="us-was-1",
-            )
+        mock_metadata = MagicMock()
+        mock_metadata.url = "https://test-sandbox.us-was-1.bl.run"
+        mock_sandbox.metadata = mock_metadata
+        mock_sandbox_instance.return_value = mock_sandbox
+        
+        # Mock kernel creation
+        mock_create_kernel.return_value = "kernel-123"
+        
+        executor = BlaxelExecutor(
+            additional_imports=["numpy"],
+            logger=logger,
+            sandbox_name="test-sandbox",
+            image="custom-image:latest",
+            memory=8192,
+            region="us-was-1",
+        )
 
-            assert executor.sandbox_name == "test-sandbox"
-            assert executor.image == "custom-image:latest"
-            assert executor.memory == 8192
-            assert executor.region == "us-was-1"
+        assert executor.sandbox_name == "test-sandbox"
+        assert executor.image == "custom-image:latest"
+        assert executor.memory == 8192
+        assert executor.region == "us-was-1"
+        assert mock_install_packages.called
 
+    @patch("smolagents.remote_executors._create_kernel_http")
     @patch("blaxel.core.SandboxInstance")
+    @patch("blaxel.core.client.api.compute.create_sandbox")
     @patch("blaxel.core.client.api.compute.delete_sandbox")
-    def test_blaxel_executor_cleanup(self, mock_delete_sandbox, mock_sandbox_instance):
+    @patch("blaxel.core.settings")
+    def test_blaxel_executor_cleanup(
+        self, mock_settings, mock_delete_sandbox, mock_create_sandbox, mock_sandbox_instance, mock_create_kernel
+    ):
         """Test BlaxelExecutor cleanup method."""
         logger = MagicMock()
+        mock_settings.headers = {}
+        
+        # Mock sandbox response
+        mock_response = MagicMock()
+        mock_create_sandbox.sync.return_value = mock_response
+        
+        # Mock SandboxInstance
         mock_sandbox = MagicMock()
-        mock_sandbox_instance.create.return_value = mock_sandbox
-
-        with patch("asyncio.run") as mock_asyncio_run:
-            mock_asyncio_run.return_value = mock_sandbox
-            executor = BlaxelExecutor(additional_imports=[], logger=logger)
-
-            # Test cleanup
-            executor.cleanup()
+        mock_metadata = MagicMock()
+        mock_metadata.url = "https://test-sandbox.bl.run"
+        mock_sandbox.metadata = mock_metadata
+        mock_sandbox_instance.return_value = mock_sandbox
+        
+        # Mock kernel creation
+        mock_create_kernel.return_value = "kernel-123"
+        
+        executor = BlaxelExecutor(additional_imports=[], logger=logger)
+
+        # Test cleanup
+        executor.cleanup()
 
-            # Verify that delete_sandbox.sync was called
-            assert mock_delete_sandbox.sync.called
-            # Verify sandbox reference was cleaned up
-            assert not hasattr(executor, "sandbox")
+        # Verify that delete_sandbox.sync was called
+        assert mock_delete_sandbox.sync.called
+        # Verify sandbox reference was cleaned up
+        assert not hasattr(executor, "sandbox")