def test_hydrate_registration_parameters__workflow_nothing_set():
workflow = _workflow_pb2.WorkflowSpec(
template=_core_workflow_pb2.WorkflowTemplate(
id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
name="name",
),
nodes=[
_core_workflow_pb2.Node(
id="foo",
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK, name="task1")),
)
],
))
identifier, entity = hydrate_registration_parameters(
_identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow)
assert (identifier == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project",
domain="domain",
name="name",
version="12345",
) == entity.template.id)
assert len(workflow.template.nodes) == 1
assert workflow.template.nodes[
0].task_node.reference_id == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK,
project="project",
domain="domain",
name="task1",
version="12345",
)
def to_flyte_idl(self):
"""
:rtype: flyteidl.core.workflow_pb2.WorkflowTemplate
"""
return _core_workflow.WorkflowTemplate(
id=self.id.to_flyte_idl(),
metadata=self.metadata.to_flyte_idl(),
interface=self.interface.to_flyte_idl(),
nodes=[n.to_flyte_idl() for n in self.nodes],
outputs=[o.to_flyte_idl() for o in self.outputs],
failure_node=self.failure_node.to_flyte_idl() if self.failure_node is not None else None
)
def test_hydrate_registration_parameters__subworkflows():
workflow_template = _core_workflow_pb2.WorkflowTemplate()
workflow_template.id.CopyFrom(
_identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW,
name="workflow"))
sub_workflow_template = _core_workflow_pb2.WorkflowTemplate()
sub_workflow_template.id.CopyFrom(
_identifier_pb2.Identifier(resource_type=_identifier_pb2.WORKFLOW,
name="subworkflow"))
sub_workflow_template.nodes.append(
_core_workflow_pb2.Node(
id="task_node",
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK)),
))
workflow_spec = _workflow_pb2.WorkflowSpec(template=workflow_template)
workflow_spec.sub_workflows.append(sub_workflow_template)
identifier, entity = hydrate_registration_parameters(
_identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow_spec)
assert (identifier == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project",
domain="domain",
name="workflow",
version="12345",
) == entity.template.id)
assert entity.sub_workflows[0].id == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project",
domain="domain",
name="subworkflow",
version="12345",
)
def get_workflow_template():
"""
This function retrieves a TasKTemplate object from the pb file in the resources directory.
It was created by reading from Flyte Admin, the following workflow, after registration.
from __future__ import absolute_import
from flytekit.common.types.primitives import Integer
from flytekit.sdk.tasks import (
python_task,
inputs,
outputs,
)
from flytekit.sdk.types import Types
from flytekit.sdk.workflow import workflow_class, Input, Output
@inputs(a=Types.Integer)
@outputs(b=Types.Integer, c=Types.Integer)
@python_task()
def demo_task_for_promote(wf_params, a, b, c):
b.set(a + 1)
c.set(a + 2)
@workflow_class()
class OneTaskWFForPromote(object):
wf_input = Input(Types.Integer, required=True)
my_task_node = demo_task_for_promote(a=wf_input)
wf_output_b = Output(my_task_node.outputs.b, sdk_type=Integer)
wf_output_c = Output(my_task_node.outputs.c, sdk_type=Integer)
:rtype: flytekit.models.core.workflow.WorkflowTemplate
"""
workflow_template_pb = _workflow_pb2.WorkflowTemplate()
# So that tests that use this work when run from any directory
basepath = _path.dirname(__file__)
filepath = _path.abspath(
_path.join(basepath, "resources/protos", "OneTaskWFForPromote.pb"))
with open(filepath, "rb") as fh:
workflow_template_pb.ParseFromString(fh.read())
wt = _workflow_model.WorkflowTemplate.from_flyte_idl(workflow_template_pb)
return wt
def test_hydrate_registration_parameters__workflow_already_set():
workflow = _workflow_pb2.WorkflowSpec(
template=_core_workflow_pb2.WorkflowTemplate(
id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project2",
domain="domain2",
name="name",
version="abc",
), ))
identifier, entity = hydrate_registration_parameters(
_identifier_pb2.WORKFLOW, "project", "domain", "12345", workflow)
assert (identifier == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project2",
domain="domain2",
name="name",
version="abc",
) == entity.template.id)
def test_hydrate_workflow_template__branch_node():
workflow_template = _core_workflow_pb2.WorkflowTemplate()
branch_node = _core_workflow_pb2.Node(
id="branch_node",
branch_node=_core_workflow_pb2.BranchNode(
if_else=_core_workflow_pb2.IfElseBlock(
case=_core_workflow_pb2.IfBlock(
then_node=_core_workflow_pb2.Node(
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK,
name="if_case"), ), )),
else_node=_core_workflow_pb2.Node(
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK,
name="else_node"), ), ),
), ),
)
def test_hydrate_workflow_template():
workflow_template = _core_workflow_pb2.WorkflowTemplate()
workflow_template.nodes.append(
_core_workflow_pb2.Node(
id="task_node",
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK)),
))
workflow_template.nodes.append(
_core_workflow_pb2.Node(
id="launchplan_ref",
workflow_node=_core_workflow_pb2.WorkflowNode(
launchplan_ref=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.LAUNCH_PLAN,
project="project2",
)),
))
workflow_template.nodes.append(
_core_workflow_pb2.Node(
id="sub_workflow_ref",
workflow_node=_core_workflow_pb2.WorkflowNode(
sub_workflow_ref=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project2",
domain="domain2",
)),
))
workflow_template.nodes.append(
_core_workflow_pb2.Node(
id="unchanged",
task_node=_core_workflow_pb2.TaskNode(
reference_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.TASK,
project="project2",
domain="domain2",
version="abc")),
))
hydrated_workflow_template = _hydrate_workflow_template_nodes(
"project", "domain", "12345", workflow_template)
assert len(hydrated_workflow_template.nodes) == 4
task_node_identifier = hydrated_workflow_template.nodes[
0].task_node.reference_id
assert task_node_identifier.project == "project"
assert task_node_identifier.domain == "domain"
assert task_node_identifier.version == "12345"
launchplan_ref_identifier = hydrated_workflow_template.nodes[
1].workflow_node.launchplan_ref
assert launchplan_ref_identifier.project == "project2"
assert launchplan_ref_identifier.domain == "domain"
assert launchplan_ref_identifier.version == "12345"
sub_workflow_ref_identifier = hydrated_workflow_template.nodes[
2].workflow_node.sub_workflow_ref
assert sub_workflow_ref_identifier.project == "project2"
assert sub_workflow_ref_identifier.domain == "domain2"
assert sub_workflow_ref_identifier.version == "12345"
unchanged_identifier = hydrated_workflow_template.nodes[
3].task_node.reference_id
assert unchanged_identifier.project == "project2"
assert unchanged_identifier.domain == "domain2"
assert unchanged_identifier.version == "abc"
