mirror of
https://github.com/langgenius/dify.git
synced 2026-07-27 15:08:35 +08:00
test: use sqlite3 session in test_workflow_node_execution_offload (#38728)
This commit is contained in:
Asuka Minato
GitHub
parent
00496ec867
commit
a45f862fb2
@ -1,212 +1,204 @@
|
||||
"""
|
||||
Unit tests for WorkflowNodeExecutionOffload model, focusing on process_data truncation functionality.
|
||||
"""
|
||||
"""SQLite-backed tests for workflow node execution offload behavior."""
|
||||
|
||||
import json
|
||||
from datetime import UTC, datetime
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from extensions.storage.storage_type import StorageType
|
||||
from graphon.enums import WorkflowNodeExecutionStatus
|
||||
from models.enums import CreatorUserRole, ExecutionOffLoadType
|
||||
from models.model import UploadFile
|
||||
from models.workflow import WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload
|
||||
from models.workflow import (
|
||||
WorkflowNodeExecutionModel,
|
||||
WorkflowNodeExecutionOffload,
|
||||
WorkflowNodeExecutionTriggeredFrom,
|
||||
)
|
||||
|
||||
TABLES = (WorkflowNodeExecutionModel, WorkflowNodeExecutionOffload, UploadFile)
|
||||
|
||||
|
||||
def _persist_upload_file(session: Session, *, key: str = "offload/process-data.json") -> UploadFile:
|
||||
upload_file = UploadFile(
|
||||
tenant_id="tenant-1",
|
||||
storage_type=StorageType.LOCAL,
|
||||
key=key,
|
||||
name="process-data.json",
|
||||
size=18,
|
||||
extension="json",
|
||||
mime_type="application/json",
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by="account-1",
|
||||
created_at=datetime.now(UTC),
|
||||
used=True,
|
||||
)
|
||||
session.add(upload_file)
|
||||
session.flush()
|
||||
return upload_file
|
||||
|
||||
|
||||
def _persist_execution(
|
||||
session: Session,
|
||||
*,
|
||||
process_data: str = '{"test": "data"}',
|
||||
offloads: tuple[tuple[ExecutionOffLoadType, str], ...] = (),
|
||||
) -> WorkflowNodeExecutionModel:
|
||||
execution = WorkflowNodeExecutionModel(
|
||||
tenant_id="tenant-1",
|
||||
app_id="app-1",
|
||||
workflow_id="workflow-1",
|
||||
triggered_from=WorkflowNodeExecutionTriggeredFrom.WORKFLOW_RUN,
|
||||
workflow_run_id="run-1",
|
||||
index=1,
|
||||
predecessor_node_id=None,
|
||||
node_execution_id="node-execution-1",
|
||||
node_id="node-1",
|
||||
node_type="code",
|
||||
title="Code",
|
||||
inputs='{"input": "value"}',
|
||||
process_data=process_data,
|
||||
outputs='{"output": "value"}',
|
||||
status=WorkflowNodeExecutionStatus.SUCCEEDED,
|
||||
error=None,
|
||||
elapsed_time=0,
|
||||
execution_metadata=None,
|
||||
created_by_role=CreatorUserRole.ACCOUNT,
|
||||
created_by="account-1",
|
||||
finished_at=None,
|
||||
)
|
||||
session.add(execution)
|
||||
session.flush()
|
||||
|
||||
session.add_all(
|
||||
[
|
||||
WorkflowNodeExecutionOffload(
|
||||
tenant_id=execution.tenant_id,
|
||||
app_id=execution.app_id,
|
||||
node_execution_id=execution.id,
|
||||
type_=type_,
|
||||
file_id=file_id,
|
||||
)
|
||||
for type_, file_id in offloads
|
||||
]
|
||||
)
|
||||
session.commit()
|
||||
session.expunge_all()
|
||||
|
||||
stmt = WorkflowNodeExecutionModel.preload_offload_data(
|
||||
select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution.id)
|
||||
)
|
||||
loaded_execution = session.scalar(stmt)
|
||||
assert loaded_execution is not None
|
||||
return loaded_execution
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sqlite_session", [TABLES], indirect=True)
|
||||
class TestWorkflowNodeExecutionModel:
|
||||
"""Test WorkflowNodeExecutionModel with process_data truncation features."""
|
||||
"""Exercise truncation flags and content restoration using persisted offload rows."""
|
||||
|
||||
def create_mock_offload_data(
|
||||
self,
|
||||
inputs_file_id: str | None = None,
|
||||
outputs_file_id: str | None = None,
|
||||
process_data_file_id: str | None = None,
|
||||
) -> WorkflowNodeExecutionOffload:
|
||||
"""Create a mock offload data object."""
|
||||
offload = Mock(spec=WorkflowNodeExecutionOffload)
|
||||
offload.inputs_file_id = inputs_file_id
|
||||
offload.outputs_file_id = outputs_file_id
|
||||
offload.process_data_file_id = process_data_file_id
|
||||
|
||||
# Mock file objects
|
||||
if inputs_file_id:
|
||||
offload.inputs_file = Mock(spec=UploadFile)
|
||||
else:
|
||||
offload.inputs_file = None
|
||||
|
||||
if outputs_file_id:
|
||||
offload.outputs_file = Mock(spec=UploadFile)
|
||||
else:
|
||||
offload.outputs_file = None
|
||||
|
||||
if process_data_file_id:
|
||||
offload.process_data_file = Mock(spec=UploadFile)
|
||||
else:
|
||||
offload.process_data_file = None
|
||||
|
||||
return offload
|
||||
|
||||
def test_process_data_truncated_property_false_when_no_offload_data(self):
|
||||
"""Test process_data_truncated returns False when no offload_data."""
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
execution.offload_data = []
|
||||
def test_process_data_truncated_property_false_when_no_offload_data(self, sqlite_session: Session):
|
||||
execution = _persist_execution(sqlite_session)
|
||||
|
||||
assert execution.process_data_truncated is False
|
||||
|
||||
def test_process_data_truncated_property_false_when_no_process_data_file(self):
|
||||
"""Test process_data_truncated returns False when no process_data file."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Create real offload instances for inputs and outputs but not process_data
|
||||
inputs_offload = WorkflowNodeExecutionOffload()
|
||||
inputs_offload.type_ = ExecutionOffLoadType.INPUTS
|
||||
inputs_offload.file_id = "inputs-file"
|
||||
|
||||
outputs_offload = WorkflowNodeExecutionOffload()
|
||||
outputs_offload.type_ = ExecutionOffLoadType.OUTPUTS
|
||||
outputs_offload.file_id = "outputs-file"
|
||||
|
||||
execution.offload_data = [inputs_offload, outputs_offload]
|
||||
def test_process_data_truncated_property_false_when_no_process_data_file(self, sqlite_session: Session):
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=(
|
||||
(ExecutionOffLoadType.INPUTS, "inputs-file"),
|
||||
(ExecutionOffLoadType.OUTPUTS, "outputs-file"),
|
||||
),
|
||||
)
|
||||
|
||||
assert execution.process_data_truncated is False
|
||||
|
||||
def test_process_data_truncated_property_true_when_process_data_file_exists(self):
|
||||
"""Test process_data_truncated returns True when process_data file exists."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Create a real offload instance for process_data
|
||||
process_data_offload = WorkflowNodeExecutionOffload()
|
||||
process_data_offload.type_ = ExecutionOffLoadType.PROCESS_DATA
|
||||
process_data_offload.file_id = "process-data-file-id"
|
||||
execution.offload_data = [process_data_offload]
|
||||
def test_process_data_truncated_property_true_when_process_data_file_exists(self, sqlite_session: Session):
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file-id"),),
|
||||
)
|
||||
|
||||
assert execution.process_data_truncated is True
|
||||
|
||||
def test_load_full_process_data_with_no_offload_data(self):
|
||||
"""Test load_full_process_data when no offload data exists."""
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
execution.offload_data = []
|
||||
execution.process_data = '{"test": "data"}'
|
||||
def test_load_full_process_data_with_no_offload_data(self, sqlite_session: Session):
|
||||
execution = _persist_execution(sqlite_session)
|
||||
storage = Mock()
|
||||
|
||||
# Mock session and storage
|
||||
mock_session = Mock()
|
||||
mock_storage = Mock()
|
||||
|
||||
result = execution.load_full_process_data(mock_session, mock_storage)
|
||||
result = execution.load_full_process_data(sqlite_session, storage)
|
||||
|
||||
assert result == {"test": "data"}
|
||||
storage.load.assert_not_called()
|
||||
|
||||
def test_load_full_process_data_with_no_file(self):
|
||||
"""Test load_full_process_data when no process_data file exists."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
def test_load_full_process_data_with_no_file(self, sqlite_session: Session):
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=((ExecutionOffLoadType.INPUTS, "inputs-file"),),
|
||||
)
|
||||
storage = Mock()
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Create offload data for inputs only, not process_data
|
||||
inputs_offload = WorkflowNodeExecutionOffload()
|
||||
inputs_offload.type_ = ExecutionOffLoadType.INPUTS
|
||||
inputs_offload.file_id = "inputs-file"
|
||||
|
||||
execution.offload_data = [inputs_offload]
|
||||
execution.process_data = '{"test": "data"}'
|
||||
|
||||
# Mock session and storage
|
||||
mock_session = Mock()
|
||||
mock_storage = Mock()
|
||||
|
||||
result = execution.load_full_process_data(mock_session, mock_storage)
|
||||
result = execution.load_full_process_data(sqlite_session, storage)
|
||||
|
||||
assert result == {"test": "data"}
|
||||
storage.load.assert_not_called()
|
||||
|
||||
def test_load_full_process_data_with_file(self):
|
||||
"""Test load_full_process_data when process_data file exists."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Create process_data offload
|
||||
process_data_offload = WorkflowNodeExecutionOffload()
|
||||
process_data_offload.type_ = ExecutionOffLoadType.PROCESS_DATA
|
||||
process_data_offload.file_id = "file-id"
|
||||
|
||||
execution.offload_data = [process_data_offload]
|
||||
execution.process_data = '{"truncated": "data"}'
|
||||
|
||||
# Mock session and storage
|
||||
mock_session = Mock()
|
||||
mock_storage = Mock()
|
||||
|
||||
# Mock the _load_full_content method to return full data
|
||||
def test_load_full_process_data_with_file(self, sqlite_session: Session):
|
||||
upload_file = _persist_upload_file(sqlite_session)
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
process_data='{"truncated": true}',
|
||||
offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),),
|
||||
)
|
||||
full_process_data = {"full": "data", "large_field": "x" * 10000}
|
||||
storage = Mock()
|
||||
storage.load.return_value = json.dumps(full_process_data)
|
||||
|
||||
with pytest.MonkeyPatch.context() as mp:
|
||||
# Mock the _load_full_content method
|
||||
def mock_load_full_content(session, file_id, storage):
|
||||
assert session == mock_session
|
||||
assert file_id == "file-id"
|
||||
assert storage == mock_storage
|
||||
return full_process_data
|
||||
result = execution.load_full_process_data(sqlite_session, storage)
|
||||
|
||||
mp.setattr(execution, "_load_full_content", mock_load_full_content)
|
||||
assert result == full_process_data
|
||||
storage.load.assert_called_once_with(upload_file.key)
|
||||
|
||||
result = execution.load_full_process_data(mock_session, mock_storage)
|
||||
def test_consistency_with_inputs_outputs_truncation(self, sqlite_session: Session):
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=(
|
||||
(ExecutionOffLoadType.INPUTS, "inputs-file"),
|
||||
(ExecutionOffLoadType.OUTPUTS, "outputs-file"),
|
||||
(ExecutionOffLoadType.PROCESS_DATA, "process-data-file"),
|
||||
),
|
||||
)
|
||||
|
||||
assert result == full_process_data
|
||||
|
||||
def test_consistency_with_inputs_outputs_truncation(self):
|
||||
"""Test that process_data truncation behaves consistently with inputs/outputs."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Create offload data for all three types
|
||||
inputs_offload = WorkflowNodeExecutionOffload()
|
||||
inputs_offload.type_ = ExecutionOffLoadType.INPUTS
|
||||
inputs_offload.file_id = "inputs-file"
|
||||
|
||||
outputs_offload = WorkflowNodeExecutionOffload()
|
||||
outputs_offload.type_ = ExecutionOffLoadType.OUTPUTS
|
||||
outputs_offload.file_id = "outputs-file"
|
||||
|
||||
process_data_offload = WorkflowNodeExecutionOffload()
|
||||
process_data_offload.type_ = ExecutionOffLoadType.PROCESS_DATA
|
||||
process_data_offload.file_id = "process-data-file"
|
||||
|
||||
execution.offload_data = [inputs_offload, outputs_offload, process_data_offload]
|
||||
|
||||
# All three should be truncated
|
||||
assert execution.inputs_truncated is True
|
||||
assert execution.outputs_truncated is True
|
||||
assert execution.process_data_truncated is True
|
||||
|
||||
def test_mixed_truncation_states(self):
|
||||
"""Test mixed states of truncation."""
|
||||
from models.enums import ExecutionOffLoadType
|
||||
|
||||
execution = WorkflowNodeExecutionModel()
|
||||
|
||||
# Only process_data is truncated
|
||||
process_data_offload = WorkflowNodeExecutionOffload()
|
||||
process_data_offload.type_ = ExecutionOffLoadType.PROCESS_DATA
|
||||
process_data_offload.file_id = "process-data-file"
|
||||
|
||||
execution.offload_data = [process_data_offload]
|
||||
def test_mixed_truncation_states(self, sqlite_session: Session):
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=((ExecutionOffLoadType.PROCESS_DATA, "process-data-file"),),
|
||||
)
|
||||
|
||||
assert execution.inputs_truncated is False
|
||||
assert execution.outputs_truncated is False
|
||||
assert execution.process_data_truncated is True
|
||||
|
||||
def test_preload_offload_data_and_files_method_exists(self):
|
||||
"""Test that the preload method includes process_data_file."""
|
||||
# This test verifies the method exists and can be called
|
||||
# The actual SQL behavior would be tested in integration tests
|
||||
from sqlalchemy import select
|
||||
def test_preload_offload_data_and_files_loads_process_data_file(self, sqlite_session: Session):
|
||||
upload_file = _persist_upload_file(sqlite_session)
|
||||
execution = _persist_execution(
|
||||
sqlite_session,
|
||||
offloads=((ExecutionOffLoadType.PROCESS_DATA, upload_file.id),),
|
||||
)
|
||||
execution_id = execution.id
|
||||
sqlite_session.expunge_all()
|
||||
|
||||
stmt = select(WorkflowNodeExecutionModel)
|
||||
stmt = WorkflowNodeExecutionModel.preload_offload_data_and_files(
|
||||
select(WorkflowNodeExecutionModel).where(WorkflowNodeExecutionModel.id == execution_id)
|
||||
)
|
||||
preloaded_execution = sqlite_session.scalar(stmt)
|
||||
|
||||
# This should not raise an exception
|
||||
preloaded_stmt = WorkflowNodeExecutionModel.preload_offload_data_and_files(stmt)
|
||||
|
||||
# The statement should be modified (different object)
|
||||
assert preloaded_stmt is not stmt
|
||||
assert preloaded_execution is not None
|
||||
assert len(preloaded_execution.offload_data) == 1
|
||||
assert preloaded_execution.offload_data[0].file is not None
|
||||
assert preloaded_execution.offload_data[0].file.key == upload_file.key
|
||||
|
||||
Loading…
Reference in New Issue
Block a user