def pop_input_from_component_spec(
component_spec: pipeline_spec_pb2.ComponentSpec,
input_name: str,
) -> None:
"""Removes an input from component spec input_definitions.
Args:
component_spec: The component spec to update in place.
input_name: The name of the input, which could be an artifact or paremeter.
"""
component_spec.input_definitions.artifacts.pop(input_name)
component_spec.input_definitions.parameters.pop(input_name)
if component_spec.input_definitions == pipeline_spec_pb2.ComponentInputsSpec(
):
component_spec.ClearField('input_definitions')
def test_update_task_inputs_spec(self, original_task_spec,
parent_component_inputs,
tasks_in_current_dag, expected_result):
pipeline_params = self.TEST_PIPELINE_PARAMS
expected_spec = pipeline_spec_pb2.PipelineTaskSpec()
json_format.ParseDict(expected_result, expected_spec)
task_spec = pipeline_spec_pb2.PipelineTaskSpec()
json_format.ParseDict(original_task_spec, task_spec)
parent_component_inputs_spec = pipeline_spec_pb2.ComponentInputsSpec()
json_format.ParseDict(parent_component_inputs,
parent_component_inputs_spec)
dsl_component_spec.update_task_inputs_spec(
task_spec, parent_component_inputs_spec, pipeline_params,
tasks_in_current_dag)
self.assertEqual(expected_spec, task_spec)
def test_fill_in_component_input_default_value(self, parameter_type,
default_value, expected):
component_spec = pipeline_spec_pb2.ComponentSpec(
input_definitions=pipeline_spec_pb2.ComponentInputsSpec(
parameters={
'input1':
pipeline_spec_pb2.ComponentInputsSpec.ParameterSpec(
parameter_type=parameter_type)
}))
pipeline_spec_builder._fill_in_component_input_default_value(
component_spec=component_spec,
input_name='input1',
default_value=default_value)
self.assertEqual(
expected,
component_spec.input_definitions.parameters['input1'].
default_value,
)