jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(workflow_cycle_manager): Removes redundant repository methods and adds caching (#22597)

Browse source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-07-18 09:26:05 +08:00 • committed by GitHub
commit b88dd17fc1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 294 additions and 535 deletions
Download patch file Download diff file Expand all files Collapse all files

42
api/tests/unit_tests/core/workflow/test_workflow_cycle_manager.py
View file

@ -80,15 +80,12 @@ def real_workflow_system_variables():
@pytest.fixture
def mock_node_execution_repository():
repo = MagicMock(spec=WorkflowNodeExecutionRepository)
repo.get_by_node_execution_id.return_value = None
repo.get_running_executions.return_value = []
return repo
@pytest.fixture
def mock_workflow_execution_repository():
repo = MagicMock(spec=WorkflowExecutionRepository)
repo.get.return_value = None
return repo
@ -217,8 +214,8 @@ def test_handle_workflow_run_success(workflow_cycle_manager, mock_workflow_execu
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Pre-populate the cache with the workflow execution
workflow_cycle_manager._workflow_execution_cache[workflow_execution.id_] = workflow_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_run_success(
@ -251,11 +248,10 @@ def test_handle_workflow_run_failed(workflow_cycle_manager, mock_workflow_execut
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Pre-populate the cache with the workflow execution
workflow_cycle_manager._workflow_execution_cache[workflow_execution.id_] = workflow_execution
# Mock get_running_executions to return an empty list
workflow_cycle_manager._workflow_node_execution_repository.get_running_executions.return_value = []
# No running node executions in cache (empty cache)
# Call the method
result = workflow_cycle_manager.handle_workflow_run_failed(
@ -289,8 +285,8 @@ def test_handle_node_execution_start(workflow_cycle_manager, mock_workflow_execu
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Pre-populate the cache with the workflow execution
workflow_cycle_manager._workflow_execution_cache[workflow_execution.id_] = workflow_execution
# Create a mock event
event = MagicMock(spec=QueueNodeStartedEvent)
@ -342,8 +338,8 @@ def test_get_workflow_execution_or_raise_error(workflow_cycle_manager, mock_work
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository get method to return the real execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Pre-populate the cache with the workflow execution
workflow_cycle_manager._workflow_execution_cache["test-workflow-run-id"] = workflow_execution
# Call the method
result = workflow_cycle_manager._get_workflow_execution_or_raise_error("test-workflow-run-id")
@ -351,11 +347,13 @@ def test_get_workflow_execution_or_raise_error(workflow_cycle_manager, mock_work
# Verify the result
assert result == workflow_execution
# Test error case
workflow_cycle_manager._workflow_execution_repository.get.return_value = None
# Test error case - clear cache
workflow_cycle_manager._workflow_execution_cache.clear()
# Expect an error when execution is not found
with pytest.raises(ValueError):
from core.app.task_pipeline.exc import WorkflowRunNotFoundError
with pytest.raises(WorkflowRunNotFoundError):
workflow_cycle_manager._get_workflow_execution_or_raise_error("non-existent-id")
@ -384,8 +382,8 @@ def test_handle_workflow_node_execution_success(workflow_cycle_manager):
created_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository to return the node execution
workflow_cycle_manager._workflow_node_execution_repository.get_by_node_execution_id.return_value = node_execution
# Pre-populate the cache with the node execution
workflow_cycle_manager._node_execution_cache["test-node-execution-id"] = node_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_node_execution_success(
@ -414,8 +412,8 @@ def test_handle_workflow_run_partial_success(workflow_cycle_manager, mock_workfl
started_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock _get_workflow_execution_or_raise_error to return the real workflow_execution
workflow_cycle_manager._workflow_execution_repository.get.return_value = workflow_execution
# Pre-populate the cache with the workflow execution
workflow_cycle_manager._workflow_execution_cache[workflow_execution.id_] = workflow_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_run_partial_success(
@ -462,8 +460,8 @@ def test_handle_workflow_node_execution_failed(workflow_cycle_manager):
created_at=datetime.now(UTC).replace(tzinfo=None),
)
# Mock the repository to return the node execution
workflow_cycle_manager._workflow_node_execution_repository.get_by_node_execution_id.return_value = node_execution
# Pre-populate the cache with the node execution
workflow_cycle_manager._node_execution_cache["test-node-execution-id"] = node_execution
# Call the method
result = workflow_cycle_manager.handle_workflow_node_execution_failed(