def test_hydrate_registration_parameters__launch_plan_already_set():
launch_plan = _launch_plan_pb2.LaunchPlanSpec(
workflow_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project2",
domain="domain2",
name="workflow_name",
version="abc",
)
)
identifier, entity = hydrate_registration_parameters(
_identifier_pb2.Identifier(
resource_type=_identifier_pb2.LAUNCH_PLAN,
project="project2",
domain="domain2",
name="workflow_name",
version="abc",
),
"project",
"domain",
"12345",
launch_plan,
)
assert identifier == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.LAUNCH_PLAN,
project="project2",
domain="domain2",
name="workflow_name",
version="abc",
)
assert entity.workflow_id == launch_plan.workflow_id
def test_hydrate_registration_parameters__launch_plan_nothing_set():
launch_plan = _launch_plan_pb2.LaunchPlan(
id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.LAUNCH_PLAN,
name="lp_name",
),
spec=_launch_plan_pb2.LaunchPlanSpec(
workflow_id=_identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
name="workflow_name",
)),
)
identifier, entity = hydrate_registration_parameters(
_identifier_pb2.LAUNCH_PLAN,
"project",
"domain",
"12345",
launch_plan,
)
assert identifier == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.LAUNCH_PLAN,
project="project",
domain="domain",
name="lp_name",
version="12345",
)
assert entity.spec.workflow_id == _identifier_pb2.Identifier(
resource_type=_identifier_pb2.WORKFLOW,
project="project",
domain="domain",
name="workflow_name",
version="12345",
)
def to_flyte_idl(self):
"""
:rtype: flyteidl.admin.launch_plan_pb2.LaunchPlanSpec
"""
return _launch_plan.LaunchPlanSpec(
workflow_id=self.workflow_id.to_flyte_idl(),
entity_metadata=self.entity_metadata.to_flyte_idl(),
default_inputs=self.default_inputs.to_flyte_idl(),
fixed_inputs=self.fixed_inputs.to_flyte_idl(),
labels=self.labels.to_flyte_idl(),
annotations=self.annotations.to_flyte_idl(),
auth_role=self.auth_role.to_flyte_idl(),
)
def to_flyte_idl(self):
"""
:rtype: flyteidl.admin.launch_plan_pb2.LaunchPlanSpec
"""
return _launch_plan.LaunchPlanSpec(
workflow_id=self.workflow_id.to_flyte_idl(),
entity_metadata=self.entity_metadata.to_flyte_idl(),
default_inputs=self.default_inputs.to_flyte_idl(),
fixed_inputs=self.fixed_inputs.to_flyte_idl(),
labels=self.labels.to_flyte_idl(),
annotations=self.annotations.to_flyte_idl(),
auth_role=self.auth_role.to_flyte_idl()
if self.auth_role else None,
raw_output_data_config=self.raw_output_data_config.to_flyte_idl(),
max_parallelism=self.max_parallelism,
security_context=self.security_context.to_flyte_idl()
if self.security_context else None,
)
def test_old_style_role():
identifier_model = identifier.Identifier(identifier.ResourceType.TASK,
"project", "domain", "name",
"version")
s = schedule.Schedule("asdf", "1 3 4 5 6 7")
launch_plan_metadata_model = launch_plan.LaunchPlanMetadata(
schedule=s, notifications=[])
v = interface.Variable(types.LiteralType(simple=types.SimpleType.BOOLEAN),
"asdf asdf asdf")
p = interface.Parameter(var=v)
parameter_map = interface.ParameterMap({"ppp": p})
fixed_inputs = literals.LiteralMap({
"a":
literals.Literal(scalar=literals.Scalar(primitive=literals.Primitive(
integer=1)))
})
labels_model = common.Labels({})
annotations_model = common.Annotations({"my": "annotation"})
raw_data_output_config = common.RawOutputDataConfig("s3://bucket")
old_role = _launch_plan_idl.Auth(
kubernetes_service_account="my:service:account")
old_style_spec = _launch_plan_idl.LaunchPlanSpec(
workflow_id=identifier_model.to_flyte_idl(),
entity_metadata=launch_plan_metadata_model.to_flyte_idl(),
default_inputs=parameter_map.to_flyte_idl(),
fixed_inputs=fixed_inputs.to_flyte_idl(),
labels=labels_model.to_flyte_idl(),
annotations=annotations_model.to_flyte_idl(),
raw_output_data_config=raw_data_output_config.to_flyte_idl(),
auth=old_role,
)
lp_spec = launch_plan.LaunchPlanSpec.from_flyte_idl(old_style_spec)
assert lp_spec.auth_role.assumable_iam_role == "my:service:account"
