|
| 1 | +""" |
| 2 | +Unit tests for WorkflowRunner._get_block_outputs. |
| 3 | +
|
| 4 | +Verifies that Workflow blocks (returning Execution objects) emit only |
| 5 | +declared outputs, not the full child execution tree via model_dump(). |
| 6 | +""" |
| 7 | + |
| 8 | +from pydantic import BaseModel, Field |
| 9 | + |
| 10 | +from workflows_mcp.engine.execution import Execution |
| 11 | +from workflows_mcp.engine.metadata import Metadata |
| 12 | +from workflows_mcp.engine.workflow_runner import WorkflowRunner |
| 13 | + |
| 14 | + |
| 15 | +class FakeBlockOutput(BaseModel): |
| 16 | + """Mimics ShellOutput/LLMCallOutput for testing.""" |
| 17 | + |
| 18 | + stdout: str = Field(default="") |
| 19 | + exit_code: int = Field(default=0) |
| 20 | + |
| 21 | + |
| 22 | +def _make_runner() -> WorkflowRunner: |
| 23 | + """Create a bare WorkflowRunner (no callback needed for _get_block_outputs).""" |
| 24 | + return WorkflowRunner() |
| 25 | + |
| 26 | + |
| 27 | +def test_get_block_outputs_none_returns_empty_dict() -> None: |
| 28 | + """None input returns empty dict.""" |
| 29 | + runner = _make_runner() |
| 30 | + assert runner._get_block_outputs(None) == {} |
| 31 | + |
| 32 | + |
| 33 | +def test_get_block_outputs_regular_block_uses_model_dump() -> None: |
| 34 | + """Regular BlockOutput (Shell/LLMCall) returns full model_dump().""" |
| 35 | + runner = _make_runner() |
| 36 | + output = FakeBlockOutput(stdout="hello world", exit_code=0) |
| 37 | + result = runner._get_block_outputs(output) |
| 38 | + |
| 39 | + assert result == {"stdout": "hello world", "exit_code": 0} |
| 40 | + |
| 41 | + |
| 42 | +def test_get_block_outputs_execution_returns_only_declared_outputs() -> None: |
| 43 | + """Execution (Workflow block) returns only declared outputs, not full tree. |
| 44 | +
|
| 45 | + This is the core fix: Workflow blocks return Execution objects. |
| 46 | + Before the fix, model_dump() serialized the ENTIRE child execution tree |
| 47 | + including all child blocks, metadata, and nested executions (~100KB+). |
| 48 | + After the fix, only the declared outputs (~few KB) are returned. |
| 49 | + """ |
| 50 | + runner = _make_runner() |
| 51 | + |
| 52 | + # Simulate a child workflow execution with nested blocks and large content |
| 53 | + child_execution = Execution( |
| 54 | + inputs={"source_url": "https://example.com/large-doc.pdf"}, |
| 55 | + outputs={ |
| 56 | + "document_text": "extracted text (trimmed for test)", |
| 57 | + "metadata_json": '{"title": "Test Document"}', |
| 58 | + }, |
| 59 | + blocks={ |
| 60 | + "extract": Execution( |
| 61 | + inputs={"url": "https://example.com/large-doc.pdf"}, |
| 62 | + outputs={"stdout": "A" * 84000}, # Simulates 84K document text |
| 63 | + metadata=Metadata.create_leaf_success( |
| 64 | + type="Shell", |
| 65 | + id="extract", |
| 66 | + duration_ms=1500, |
| 67 | + started_at="2026-01-01T00:00:00Z", |
| 68 | + wave=0, |
| 69 | + execution_order=0, |
| 70 | + index=0, |
| 71 | + depth=1, |
| 72 | + ), |
| 73 | + ), |
| 74 | + "metadata": Execution( |
| 75 | + inputs={"text": "some text"}, |
| 76 | + outputs={"response": '{"title": "Test Document"}'}, |
| 77 | + metadata=Metadata.create_leaf_success( |
| 78 | + type="LLMCall", |
| 79 | + id="metadata", |
| 80 | + duration_ms=2000, |
| 81 | + started_at="2026-01-01T00:00:01Z", |
| 82 | + wave=1, |
| 83 | + execution_order=1, |
| 84 | + index=1, |
| 85 | + depth=1, |
| 86 | + ), |
| 87 | + ), |
| 88 | + }, |
| 89 | + ) |
| 90 | + |
| 91 | + result = runner._get_block_outputs(child_execution) |
| 92 | + |
| 93 | + # Should return ONLY declared outputs |
| 94 | + assert result == { |
| 95 | + "document_text": "extracted text (trimmed for test)", |
| 96 | + "metadata_json": '{"title": "Test Document"}', |
| 97 | + } |
| 98 | + |
| 99 | + # Crucially: should NOT contain the full child block tree |
| 100 | + assert "blocks" not in result |
| 101 | + assert "inputs" not in result |
| 102 | + assert "metadata" not in result |
| 103 | + # And should NOT contain the 84K content from child blocks |
| 104 | + assert "A" * 84000 not in str(result) |
| 105 | + |
| 106 | + |
| 107 | +def test_get_block_outputs_execution_empty_outputs() -> None: |
| 108 | + """Execution with empty outputs returns empty dict.""" |
| 109 | + runner = _make_runner() |
| 110 | + |
| 111 | + child_execution = Execution( |
| 112 | + inputs={"key": "value"}, |
| 113 | + outputs={}, |
| 114 | + blocks={"some_block": Execution(inputs={}, outputs={"data": "big"})}, |
| 115 | + ) |
| 116 | + |
| 117 | + result = runner._get_block_outputs(child_execution) |
| 118 | + assert result == {} |
| 119 | + |
| 120 | + |
| 121 | +def test_get_block_outputs_does_not_affect_regular_block_with_extra_fields() -> None: |
| 122 | + """Regular BlockOutput with extra fields (custom outputs) still works.""" |
| 123 | + |
| 124 | + class CustomOutput(BaseModel): |
| 125 | + model_config = {"extra": "allow"} |
| 126 | + stdout: str = "" |
| 127 | + exit_code: int = 0 |
| 128 | + |
| 129 | + runner = _make_runner() |
| 130 | + output = CustomOutput(stdout="ok", exit_code=0, custom_field="custom_value") |
| 131 | + result = runner._get_block_outputs(output) |
| 132 | + |
| 133 | + assert result["stdout"] == "ok" |
| 134 | + assert result["exit_code"] == 0 |
| 135 | + assert result["custom_field"] == "custom_value" |
0 commit comments