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.TaskNode
"""
return _core_workflow.TaskNode(
reference_id=self.reference_id.to_flyte_idl(),
overrides=self.overrides.to_flyte_idl()
if self.overrides is not None else None,
)
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_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 to_flyte_idl(self):
"""
:rtype: flyteidl.core.workflow_pb2.TaskNode
"""
return _core_workflow.TaskNode(reference_id=self.reference_id.to_flyte_idl())
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"
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"), ), ),
), ),
)
branch_node.branch_node.if_else.other.extend([
_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="other_1"), ), ), ),
_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="other_2"), ), ), ),
])
workflow_template.nodes.append(branch_node)
hydrated_workflow_template = _hydrate_workflow_template_nodes(
"project", "domain", "12345", workflow_template)
if_case_id = hydrated_workflow_template.nodes[
0].branch_node.if_else.case.then_node.task_node.reference_id
assert if_case_id.project == "project"
assert if_case_id.domain == "domain"
assert if_case_id.name == "if_case"
