def test_get_node_execution_outputs(mock_client_factory, execution_data_locations): mock_client = MagicMock() mock_client.get_node_execution_data = MagicMock( return_value=_execution_models.NodeExecutionGetDataResponse( execution_data_locations[0], execution_data_locations[1] ) ) mock_client_factory.return_value = mock_client m = MagicMock() type(m).id = PropertyMock( return_value=identifier.NodeExecutionIdentifier( "node-a", identifier.WorkflowExecutionIdentifier( "project", "domain", "name", ) ) ) inputs = engine.FlyteNodeExecution(m).get_outputs() assert len(inputs.literals) == 1 assert inputs.literals['b'].scalar.primitive.integer == 2 mock_client.get_node_execution_data.assert_called_once_with( identifier.NodeExecutionIdentifier( "node-a", identifier.WorkflowExecutionIdentifier( "project", "domain", "name", ) ) )
def test_get_full_node_execution_outputs(mock_client_factory): mock_client = MagicMock() mock_client.get_node_execution_data = MagicMock( return_value=_execution_models.NodeExecutionGetDataResponse(None, None, _INPUT_MAP, _OUTPUT_MAP) ) mock_client_factory.return_value = mock_client m = MagicMock() type(m).id = PropertyMock( return_value=identifier.NodeExecutionIdentifier( "node-a", identifier.WorkflowExecutionIdentifier( "project", "domain", "name", ), ) ) outputs = engine.FlyteNodeExecution(m).get_outputs() assert len(outputs.literals) == 1 assert outputs.literals["b"].scalar.primitive.integer == 2 mock_client.get_node_execution_data.assert_called_once_with( identifier.NodeExecutionIdentifier( "node-a", identifier.WorkflowExecutionIdentifier( "project", "domain", "name", ), ) )