From a45f862fb29ca26219e93906f1f8886344efdbcc Mon Sep 17 00:00:00 2001 From: Asuka Minato Date: Sat, 25 Jul 2026 23:23:32 +0900 Subject: [PATCH] test: use sqlite3 session in test_workflow_node_execution_offload (#38728) --- .../test_workflow_node_execution_offload.py | 330 +++++++++--------- 1 file changed, 161 insertions(+), 169 deletions(-) diff --git a/api/tests/unit_tests/models/test_workflow_node_execution_offload.py b/api/tests/unit_tests/models/test_workflow_node_execution_offload.py index c5fd6511dfd..416ba10ce37 100644 --- a/api/tests/unit_tests/models/test_workflow_node_execution_offload.py +++ b/api/tests/unit_tests/models/test_workflow_node_execution_offload.py @@ -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