Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 65 additions & 0 deletions ee/vellum_ee/workflows/display/tests/test_mocks.py
Original file line number Diff line number Diff line change
Expand Up @@ -779,3 +779,68 @@ class Outputs(BaseWorkflow.Outputs):

# AND the then_outputs should be serialized correctly
assert serialized_mock["then_outputs"]["result"] == "mocked_result"


def test_mocks__validate_all__try_node_wrapped_node_deserializes_correctly():
"""
Tests that validate_all correctly deserializes a mock for a node wrapped in a TryNode adornment.
"""

# GIVEN a node wrapped in a TryNode adornment
@TryNode.wrap()
class WrappedNode(BaseNode):
class Outputs(BaseNode.Outputs):
result: str

# AND a workflow that uses the wrapped node
class MyWorkflow(BaseWorkflow):
graph = WrappedNode

class Outputs(BaseWorkflow.Outputs):
final_result = WrappedNode.Outputs.result

# AND the inner wrapped node's ID
inner_node = WrappedNode.__wrapped_node__
assert inner_node is not None

# AND a raw mock workflow node execution using the inner node's ID
raw_mock_workflow_node_executions = [
{
"node_id": str(inner_node.__id__),
"when_condition": {
"type": "BINARY_EXPRESSION",
"operator": ">=",
"lhs": {
"type": "EXECUTION_COUNTER",
"node_id": str(inner_node.__id__),
},
"rhs": {
"type": "CONSTANT_VALUE",
"value": {
"type": "NUMBER",
"value": 1,
},
},
},
"then_outputs": {
"result": "mocked_result",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DevinAI what happens if we add an "error" output here and change nothing else about the test?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I tested this - adding an "error" output to then_outputs causes a WorkflowInitializationException with the message "Failed to validate mock node executions".

The root cause is a NodeException: Unexpected outputs: {'error'} raised from BaseOutputs.__init__ because the inner wrapped node's Outputs class only declares result: str, not error.

The validation correctly rejects outputs that aren't declared on the node's Outputs class.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Create a second test that exhibits this.

In that test, ensure that we pass. Doing so will require us updating the business logic of deserializing node outputs to ignore any outputs that are not defined in the class

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

DevinAI apply the feedback I just left in this comment

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done! I've added a second test (test_mocks__validate_all__ignores_undefined_outputs) that includes an "error" output in then_outputs and verifies it's ignored.

I also updated BaseOutputs.__get_pydantic_core_schema__ to filter out undefined outputs during deserialization. The strict validation in BaseOutputs.__init__ is preserved for normal runtime construction.

},
},
]

# WHEN we parse the raw data on MockNodeExecution
node_output_mocks = MockNodeExecution.validate_all(
raw_mock_workflow_node_executions,
MyWorkflow,
descriptor_validator=base_descriptor_validator,
)

# THEN we get a list with one MockNodeExecution object
assert node_output_mocks is not None
assert len(node_output_mocks) == 1

# AND the MockNodeExecution has the correct when_condition
assert node_output_mocks[0].when_condition == inner_node.Execution.count.greater_than_or_equal_to(1)

# AND the then_outputs is the correct type with the expected value
assert node_output_mocks[0].then_outputs == inner_node.Outputs(result="mocked_result")