youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-05-03 08:58:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(sandbox): redesign sandbox_layer & reorganize import paths

Browse Source
This commit is contained in:
Harry
2026-01-15 13:22:46 +08:00
parent c1c8b6af44
commit 63b3e71909
21 changed files with 134 additions and 179 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
api/tests/unit_tests/core/app/layers/test_sandbox_layer.py
View File

@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
import pytest
from core.app.layers.sandbox_layer import SandboxInitializationError, SandboxLayer
from core.sandbox.manager import SandboxManager
from core.sandbox import SandboxManager
from core.virtual_environment.__base.entities import Arch
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment
from core.workflow.graph_engine.layers.base import GraphEngineLayerNotInitializedError
@ -29,37 +29,6 @@ class MockVirtualEnvironment:
self._released = True
class MockSystemVariableView:
def __init__(
self,
workflow_execution_id: str | None = "test-workflow-exec-id",
app_id: str | None = "test-app-id",
):
self._workflow_execution_id = workflow_execution_id
self._app_id = app_id
@property
def workflow_execution_id(self) -> str | None:
return self._workflow_execution_id
@property
def app_id(self) -> str | None:
return self._app_id
class MockReadOnlyGraphRuntimeStateWrapper:
def __init__(
self,
workflow_execution_id: str | None = "test-workflow-exec-id",
app_id: str | None = "test-app-id",
):
self._system_variable = MockSystemVariableView(workflow_execution_id, app_id)
@property
def system_variable(self) -> MockSystemVariableView:
return self._system_variable
class MockVMBuilder:
def __init__(self, sandbox: VirtualEnvironment):
self._sandbox = sandbox
@ -81,30 +50,40 @@ def clean_sandbox_manager():
SandboxManager.clear()
@pytest.fixture
def mock_archive_storage():
with patch("core.app.layers.sandbox_layer.ArchiveSandboxStorage") as mock_class:
mock_instance = MagicMock()
mock_instance.mount.return_value = False
mock_instance.unmount.return_value = True
mock_class.return_value = mock_instance
yield mock_instance
def create_mock_builder(sandbox):
return MockVMBuilder(sandbox)
class TestSandboxLayer:
def test_init_with_parameters(self):
layer = SandboxLayer(tenant_id="test-tenant")
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="test-sandbox")
assert layer._tenant_id == "test-tenant" # pyright: ignore[reportPrivateUsage]
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
assert layer._app_id == "test-app" # pyright: ignore[reportPrivateUsage]
assert layer._sandbox_id == "test-sandbox" # pyright: ignore[reportPrivateUsage]
def test_sandbox_property_raises_when_not_initialized(self):
layer = SandboxLayer(tenant_id="test-tenant")
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="test-sandbox")
with pytest.raises(RuntimeError) as exc_info:
_ = layer.sandbox
assert "Sandbox not initialized" in str(exc_info.value)
assert "Sandbox not found" in str(exc_info.value)
def test_sandbox_property_returns_sandbox_after_initialization(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_sandbox_property_returns_sandbox_after_initialization(self, mock_archive_storage):
sandbox_id = "test-exec-id"
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id=sandbox_id)
mock_sandbox = MockVirtualEnvironment()
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper("test-exec-id")
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -114,11 +93,10 @@ class TestSandboxLayer:
assert layer.sandbox is mock_sandbox
def test_on_graph_start_creates_sandbox_and_registers_with_manager(self):
layer = SandboxLayer(tenant_id="test-tenant-123")
def test_on_graph_start_creates_sandbox_and_registers_with_manager(self, mock_archive_storage):
sandbox_id = "test-exec-123"
layer = SandboxLayer(tenant_id="test-tenant-123", app_id="test-app-123", sandbox_id=sandbox_id)
mock_sandbox = MockVirtualEnvironment()
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper("test-exec-123", "test-app-123")
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -127,12 +105,10 @@ class TestSandboxLayer:
layer.on_graph_start()
mock_create.assert_called_once_with("test-tenant-123")
assert SandboxManager.get("test-exec-123") is mock_sandbox
assert SandboxManager.get(sandbox_id) is mock_sandbox
def test_on_graph_start_raises_sandbox_initialization_error_on_failure(self):
layer = SandboxLayer(tenant_id="test-tenant")
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper("test-exec-id")
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="test-sandbox")
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -144,30 +120,18 @@ class TestSandboxLayer:
assert "Failed to initialize sandbox" in str(exc_info.value)
assert "Sandbox provider not available" in str(exc_info.value)
def test_on_graph_start_raises_when_workflow_execution_id_not_set(self):
layer = SandboxLayer(tenant_id="test-tenant")
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id=None)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with pytest.raises(RuntimeError) as exc_info:
layer.on_graph_start()
assert "workflow_execution_id is not set" in str(exc_info.value)
def test_on_event_is_noop(self):
layer = SandboxLayer(tenant_id="test-tenant")
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="test-sandbox")
layer.on_event(GraphRunStartedEvent())
layer.on_event(GraphRunSucceededEvent(outputs={}))
layer.on_event(GraphRunFailedEvent(error="test error", exceptions_count=1))
def test_on_graph_end_releases_sandbox_and_unregisters_from_manager(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_on_graph_end_releases_sandbox_and_unregisters_from_manager(self, mock_archive_storage):
sandbox_id = "test-exec-456"
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata()
workflow_execution_id = "test-exec-456"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -175,21 +139,18 @@ class TestSandboxLayer:
):
layer.on_graph_start()
assert SandboxManager.has(workflow_execution_id)
assert SandboxManager.has(sandbox_id)
layer.on_graph_end(error=None)
mock_sandbox.release_environment.assert_called_once()
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
assert not SandboxManager.has(workflow_execution_id)
assert not SandboxManager.has(sandbox_id)
def test_on_graph_end_releases_sandbox_even_on_error(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_on_graph_end_releases_sandbox_even_on_error(self, mock_archive_storage):
sandbox_id = "test-exec-789"
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata()
workflow_execution_id = "test-exec-789"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -200,17 +161,14 @@ class TestSandboxLayer:
layer.on_graph_end(error=Exception("Workflow failed"))
mock_sandbox.release_environment.assert_called_once()
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
assert not SandboxManager.has(workflow_execution_id)
assert not SandboxManager.has(sandbox_id)
def test_on_graph_end_handles_release_failure_gracefully(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_on_graph_end_handles_release_failure_gracefully(self, mock_archive_storage):
sandbox_id = "test-exec-fail"
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata()
mock_sandbox.release_environment.side_effect = Exception("Container already removed")
workflow_execution_id = "test-exec-fail"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -221,22 +179,17 @@ class TestSandboxLayer:
layer.on_graph_end(error=None)
mock_sandbox.release_environment.assert_called_once()
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_noop_when_sandbox_not_initialized(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_on_graph_end_noop_when_sandbox_not_registered(self):
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="nonexistent-sandbox")
layer.on_graph_end(error=None)
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_is_idempotent(self):
layer = SandboxLayer(tenant_id="test-tenant")
def test_on_graph_end_is_idempotent(self, mock_archive_storage):
sandbox_id = "test-exec-idempotent"
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata()
workflow_execution_id = "test-exec-idempotent"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
with patch(
"services.sandbox.sandbox_provider_service.SandboxProviderService.create_sandbox_builder",
@ -250,7 +203,7 @@ class TestSandboxLayer:
mock_sandbox.release_environment.assert_called_once()
def test_layer_inherits_from_graph_engine_layer(self):
layer = SandboxLayer(tenant_id="test-tenant")
layer = SandboxLayer(tenant_id="test-tenant", app_id="test-app", sandbox_id="test-sandbox")
with pytest.raises(GraphEngineLayerNotInitializedError):
_ = layer.graph_runtime_state
@ -259,11 +212,9 @@ class TestSandboxLayer:
class TestSandboxLayerIntegration:
def test_full_lifecycle_with_mocked_provider(self):
layer = SandboxLayer(tenant_id="integration-tenant")
workflow_execution_id = "integration-test-exec"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
def test_full_lifecycle_with_mocked_provider(self, mock_archive_storage):
sandbox_id = "integration-test-exec"
layer = SandboxLayer(tenant_id="integration-tenant", app_id="integration-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata(sandbox_id="integration-sandbox")
@ -273,21 +224,17 @@ class TestSandboxLayerIntegration:
):
layer.on_graph_start()
assert layer._workflow_execution_id == workflow_execution_id # pyright: ignore[reportPrivateUsage]
assert layer.sandbox is mock_sandbox
assert SandboxManager.get(workflow_execution_id) is mock_sandbox
assert SandboxManager.get(sandbox_id) is mock_sandbox
layer.on_graph_end(error=None)
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
assert not SandboxManager.has(workflow_execution_id)
assert not SandboxManager.has(sandbox_id)
mock_sandbox.release_environment.assert_called_once()
def test_lifecycle_with_workflow_error(self):
layer = SandboxLayer(tenant_id="error-tenant")
workflow_execution_id = "integration-error-test"
mock_runtime_state = MockReadOnlyGraphRuntimeStateWrapper(workflow_execution_id)
layer._graph_runtime_state = mock_runtime_state # type: ignore[assignment]
def test_lifecycle_with_workflow_error(self, mock_archive_storage):
sandbox_id = "integration-error-test"
layer = SandboxLayer(tenant_id="error-tenant", app_id="error-app", sandbox_id=sandbox_id)
mock_sandbox = MagicMock(spec=VirtualEnvironment)
mock_sandbox.metadata = MockMetadata()
@ -301,6 +248,5 @@ class TestSandboxLayerIntegration:
layer.on_graph_end(error=Exception("Workflow execution failed"))
assert layer._workflow_execution_id is None # pyright: ignore[reportPrivateUsage]
assert not SandboxManager.has(workflow_execution_id)
assert not SandboxManager.has(sandbox_id)
mock_sandbox.release_environment.assert_called_once()

2
api/tests/unit_tests/core/virtual_environment/test_factory.py
View File

@ -3,7 +3,7 @@ from unittest.mock import MagicMock, patch
import pytest
from core.sandbox.factory import VMBuilder, VMType
from core.sandbox import VMBuilder, VMType
from core.virtual_environment.__base.virtual_environment import VirtualEnvironment

2
api/tests/unit_tests/core/virtual_environment/test_sandbox_manager.py
View File

@ -5,7 +5,7 @@ from typing import Any
import pytest
from core.sandbox.manager import SandboxManager
from core.sandbox import SandboxManager
from core.virtual_environment.__base.entities import (
Arch,
CommandStatus,

2
api/tests/unit_tests/core/workflow/nodes/command/test_command_node.py
View File

@ -5,7 +5,7 @@ from typing import Any
import pytest
from core.sandbox.manager import SandboxManager
from core.sandbox import SandboxManager
from core.virtual_environment.__base.entities import (
Arch,
CommandStatus,