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

feat: add tenant_id support to Sandbox and VirtualEnvironment initialization

Browse Source
This commit is contained in:
Yeuoly
2026-01-08 16:19:29 +08:00
parent 94dbda503f
commit b09a831d15
8 changed files with 66 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -48,10 +48,10 @@ class TestSandboxLayer:
"""Test SandboxLayer initialization with default parameters."""
layer = SandboxLayer()
assert layer._sandbox_type == SandboxType.DOCKER
assert layer._options == {}
assert layer._environments == {}
assert layer._sandbox is None
assert layer._sandbox_type is None # pyright: ignore[reportPrivateUsage]
assert layer._options == {} # pyright: ignore[reportPrivateUsage]
assert layer._environments == {} # pyright: ignore[reportPrivateUsage]
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_init_with_custom_parameters(self):
"""Test SandboxLayer initialization with custom parameters."""
@ -61,9 +61,9 @@ class TestSandboxLayer:
environments={"PYTHONUNBUFFERED": "1"},
)
assert layer._sandbox_type == SandboxType.LOCAL
assert layer._options == {"base_working_path": "/tmp/sandbox"}
assert layer._environments == {"PYTHONUNBUFFERED": "1"}
assert layer._sandbox_type == SandboxType.LOCAL # pyright: ignore[reportPrivateUsage]
assert layer._options == {"base_working_path": "/tmp/sandbox"} # pyright: ignore[reportPrivateUsage]
assert layer._environments == {"PYTHONUNBUFFERED": "1"} # pyright: ignore[reportPrivateUsage]
def test_sandbox_property_raises_when_not_initialized(self):
"""Test that accessing sandbox property raises error before initialization."""
@ -97,6 +97,7 @@ class TestSandboxLayer:
layer.on_graph_start()
mock_create.assert_called_once_with(
tenant_id="default",
sandbox_type=SandboxType.DOCKER,
options={"docker_image": "python:3.11"},
environments={"PATH": "/usr/bin"},
@ -110,7 +111,7 @@ class TestSandboxLayer:
with pytest.raises(SandboxInitializationError) as exc_info:
layer.on_graph_start()
assert "Failed to initialize docker sandbox" in str(exc_info.value)
assert "Failed to initialize sandbox" in str(exc_info.value)
assert "Docker not available" in str(exc_info.value)
def test_on_event_is_noop(self):
@ -134,7 +135,7 @@ class TestSandboxLayer:
layer.on_graph_end(error=None)
mock_sandbox.release_environment.assert_called_once()
assert layer._sandbox is None
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_releases_sandbox_even_on_error(self):
"""Test that on_graph_end releases sandbox even when workflow had an error."""
@ -148,7 +149,7 @@ class TestSandboxLayer:
layer.on_graph_end(error=Exception("Workflow failed"))
mock_sandbox.release_environment.assert_called_once()
assert layer._sandbox is None
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_handles_release_failure_gracefully(self):
"""Test that on_graph_end handles release failures without raising."""
@ -164,7 +165,7 @@ class TestSandboxLayer:
layer.on_graph_end(error=None)
mock_sandbox.release_environment.assert_called_once()
assert layer._sandbox is None
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_noop_when_sandbox_not_initialized(self):
"""Test that on_graph_end is a no-op when sandbox was never initialized."""
@ -173,7 +174,7 @@ class TestSandboxLayer:
# Should not raise exception
layer.on_graph_end(error=None)
assert layer._sandbox is None
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_on_graph_end_is_idempotent(self):
"""Test that calling on_graph_end multiple times is safe."""
@ -215,7 +216,7 @@ class TestSandboxLayerIntegration:
layer.on_graph_start()
# Verify sandbox is created
assert layer._sandbox is not None
assert layer._sandbox is not None # pyright: ignore[reportPrivateUsage]
sandbox_id = layer.sandbox.metadata.id
assert sandbox_id is not None
@ -223,7 +224,7 @@ class TestSandboxLayerIntegration:
layer.on_graph_end(error=None)
# Verify sandbox is released
assert layer._sandbox is None
assert layer._sandbox is None # pyright: ignore[reportPrivateUsage]
def test_lifecycle_with_workflow_error(self, tmp_path: Path):
"""Test lifecycle when workflow encounters an error."""

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

@ -40,14 +40,17 @@ class TestSandboxFactory:
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = SandboxFactory.create(
tenant_id="test-tenant",
sandbox_type=SandboxType.DOCKER,
options={"docker_image": "python:3.11-slim"},
environments={"PYTHONUNBUFFERED": "1"},
)
mock_sandbox_class.assert_called_once_with(
tenant_id="test-tenant",
options={"docker_image": "python:3.11-slim"},
environments={"PYTHONUNBUFFERED": "1"},
user_id=None,
)
assert result is mock_sandbox_instance
@ -57,9 +60,13 @@ class TestSandboxFactory:
mock_sandbox_class = MagicMock(return_value=mock_sandbox_instance)
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
SandboxFactory.create(sandbox_type=SandboxType.DOCKER, options=None, environments=None)
SandboxFactory.create(
tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER, options=None, environments=None
)
mock_sandbox_class.assert_called_once_with(options={}, environments={})
mock_sandbox_class.assert_called_once_with(
tenant_id="test-tenant", options={}, environments={}, user_id=None
)
def test_create_with_default_parameters(self):
"""Test sandbox creation with default parameters."""
@ -67,9 +74,11 @@ class TestSandboxFactory:
mock_sandbox_class = MagicMock(return_value=mock_sandbox_instance)
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
result = SandboxFactory.create(sandbox_type=SandboxType.DOCKER)
result = SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
mock_sandbox_class.assert_called_once_with(options={}, environments={})
mock_sandbox_class.assert_called_once_with(
tenant_id="test-tenant", options={}, environments={}, user_id=None
)
assert result is mock_sandbox_instance
def test_get_sandbox_class_docker_returns_correct_class(self):
@ -81,7 +90,7 @@ class TestSandboxFactory:
"core.virtual_environment.providers.docker_daemon_sandbox.DockerDaemonEnvironment",
return_value=mock_instance,
) as mock_docker_class:
SandboxFactory.create(sandbox_type=SandboxType.DOCKER)
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
mock_docker_class.assert_called_once()
def test_get_sandbox_class_local_returns_correct_class(self):
@ -92,7 +101,7 @@ class TestSandboxFactory:
"core.virtual_environment.providers.local_without_isolation.LocalVirtualEnvironment",
return_value=mock_instance,
) as mock_local_class:
SandboxFactory.create(sandbox_type=SandboxType.LOCAL)
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.LOCAL)
mock_local_class.assert_called_once()
def test_get_sandbox_class_e2b_returns_correct_class(self):
@ -103,13 +112,13 @@ class TestSandboxFactory:
"core.virtual_environment.providers.e2b_sandbox.E2BEnvironment",
return_value=mock_instance,
) as mock_e2b_class:
SandboxFactory.create(sandbox_type=SandboxType.E2B)
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.E2B)
mock_e2b_class.assert_called_once()
def test_create_with_unsupported_type_raises_value_error(self):
"""Test that unsupported sandbox type raises ValueError."""
with pytest.raises(ValueError) as exc_info:
SandboxFactory.create(sandbox_type="unsupported_type") # type: ignore[arg-type]
SandboxFactory.create(tenant_id="test-tenant", sandbox_type="unsupported_type") # type: ignore[arg-type]
assert "Unsupported sandbox type: unsupported_type" in str(exc_info.value)
@ -120,7 +129,7 @@ class TestSandboxFactory:
with patch.object(SandboxFactory, "_get_sandbox_class", return_value=mock_sandbox_class):
with pytest.raises(Exception) as exc_info:
SandboxFactory.create(sandbox_type=SandboxType.DOCKER)
SandboxFactory.create(tenant_id="test-tenant", sandbox_type=SandboxType.DOCKER)
assert "Docker daemon not available" in str(exc_info.value)
@ -131,6 +140,7 @@ class TestSandboxFactoryIntegration:
def test_create_local_sandbox_integration(self, tmp_path: Path):
"""Test creating a real local sandbox."""
sandbox = SandboxFactory.create(
tenant_id="test-tenant",
sandbox_type=SandboxType.LOCAL,
options={"base_working_path": str(tmp_path)},
environments={},

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

@ -25,7 +25,7 @@ def _drain_transport(transport: TransportReadCloser) -> bytes:
@pytest.fixture
def local_env(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> LocalVirtualEnvironment:
monkeypatch.setattr(local_without_isolation, "machine", lambda: "x86_64")
return LocalVirtualEnvironment({"base_working_path": str(tmp_path)})
return LocalVirtualEnvironment(tenant_id="test-tenant", options={"base_working_path": str(tmp_path)})
def test_construct_environment_creates_working_path(local_env: LocalVirtualEnvironment):

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

@ -28,7 +28,7 @@ class FakeSandbox(VirtualEnvironment):
self._close_streams = close_streams
self.last_execute_command: list[str] | None = None
self.released_connections: list[str] = []
super().__init__(options={}, environments={})
super().__init__(tenant_id="test-tenant", options={}, environments={})
def _construct_environment(self, options, environments): # type: ignore[override]
return Metadata(id="fake", arch=Arch.ARM64)
@ -75,6 +75,10 @@ class FakeSandbox(VirtualEnvironment):
return self._statuses.pop(0)
return CommandStatus(status=CommandStatus.Status.COMPLETED, exit_code=0)
@classmethod
def validate(cls, options: Any) -> None:
pass
def _make_node(*, command: str, working_directory: str = "") -> CommandNode:
variable_pool = VariablePool(system_variables=SystemVariable.empty(), user_inputs={})