# Copyright 2021-2022 The Kubeflow Authors # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Tests for kfp.dsl.structures.""" import os import tempfile import textwrap import unittest from absl.testing import parameterized from kfp import compiler from kfp import components from kfp import dsl from kfp.dsl import component_factory from kfp.dsl import placeholders from kfp.dsl import structures V1_YAML_IF_PLACEHOLDER = textwrap.dedent("""\ implementation: container: args: - if: cond: isPresent: optional_input_1 else: - --arg2 - default then: - --arg1 - {inputUri: optional_input_1} image: alpine inputs: - {name: optional_input_1, optional: true, type: String} name: component_if """) COMPONENT_SPEC_IF_PLACEHOLDER = structures.ComponentSpec( name='component_if', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', args=[ placeholders.IfPresentPlaceholder( input_name='optional_input_1', then=[ '--arg1', placeholders.InputUriPlaceholder( input_name='optional_input_1'), ], else_=[ '--arg2', 'default', ]) ])), inputs={ 'optional_input_1': structures.InputSpec(type='String', default=None) }, ) V1_YAML_CONCAT_PLACEHOLDER = textwrap.dedent("""\ name: component_concat implementation: container: args: - concat: ['--arg1', {inputValue: input_prefix}] image: alpine inputs: - {name: input_prefix, type: String} """) COMPONENT_SPEC_CONCAT_PLACEHOLDER = structures.ComponentSpec( name='component_concat', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', args=[ placeholders.ConcatPlaceholder(items=[ '--arg1', placeholders.InputValuePlaceholder( input_name='input_prefix'), ]) ])), inputs={'input_prefix': structures.InputSpec(type='String')}, ) V1_YAML_NESTED_PLACEHOLDER = textwrap.dedent("""\ name: component_nested implementation: container: args: - concat: - --arg1 - if: cond: isPresent: input_prefix then: {inputValue: input_prefix} else: default image: alpine inputs: - {name: input_prefix, optional: false, type: String} """) COMPONENT_SPEC_NESTED_PLACEHOLDER = structures.ComponentSpec( name='component_nested', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', args=[ placeholders.ConcatPlaceholder(items=[ '--arg1', placeholders.IfPresentPlaceholder( input_name='input_prefix', then=placeholders.InputValuePlaceholder( input_name='input_prefix'), else_='default'), ]), ])), inputs={'input_prefix': structures.InputSpec(type='String')}, ) V1_YAML_EXECUTOR_INPUT_PLACEHOLDER = textwrap.dedent("""\ name: component_executor_input inputs: - {name: input, type: String} implementation: container: image: alpine command: - python - -m - kfp.containers.entrypoint args: - --executor_input - {executorInput: null} - --function_name - test_function """) COMPONENT_SPEC_EXECUTOR_INPUT_PLACEHOLDER = structures.ComponentSpec( name='component_executor_input', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'python', '-m', 'kfp.containers.entrypoint', ], args=[ '--executor_input', placeholders.ExecutorInputPlaceholder(), '--function_name', 'test_function', ])), inputs={'input': structures.InputSpec(type='String')}, ) V1_NONCANONICAL_GENERIC_TYPES_COMPONENT_SPEC = textwrap.dedent("""\ name: generic_types_test inputs: - {name: input1, type: "List[str]"} - {name: input2, type: "typing.List[str]"} - {name: input3, type: "Dict[str, str]"} - {name: input4, type: "typing.Dict[str, str]"} outputs: - {name: output1, type: "List[str]"} - {name: output2, type: "typing.List[str]"} - {name: output3, type: "Dict[str, str]"} - {name: output4, type: "typing.Dict[str, str]"} implementation: container: image: alpine """) class StructuresTest(parameterized.TestCase): def test_component_spec_with_placeholder_referencing_nonexisting_input_output( self): with self.assertRaisesRegex( ValueError, r'^Argument "InputValuePlaceholder" references nonexistant input: "input000".' ): structures.ComponentSpec( name='component_1', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'sh', '-c', 'set -ex\necho "$0" > "$1"', placeholders.InputValuePlaceholder( input_name='input000'), placeholders.OutputPathPlaceholder( output_name='output1'), ], )), inputs={'input1': structures.InputSpec(type='String')}, outputs={'output1': structures.OutputSpec(type='String')}, ) with self.assertRaisesRegex( ValueError, r'^Argument "OutputPathPlaceholder" references nonexistant output: "output000".' ): structures.ComponentSpec( name='component_1', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'sh', '-c', 'set -ex\necho "$0" > "$1"', placeholders.InputValuePlaceholder( input_name='input1'), placeholders.OutputPathPlaceholder( output_name='output000'), ], )), inputs={'input1': structures.InputSpec(type='String')}, outputs={'output1': structures.OutputSpec(type='String')}, ) def test_simple_component_spec_save_to_component_yaml(self): # tests writing old style (less verbose) and reading in new style (more verbose) original_component_spec = structures.ComponentSpec( name='component_1', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'sh', '-c', 'set -ex\necho "$0" > "$1"', "{{$.inputs.parameters['input1']}}", "{{$.outputs.parameters['output1'].output_file}}" ], )), inputs={'input1': structures.InputSpec(type='String')}, outputs={'output1': structures.OutputSpec(type='String')}, ) test_component = components.PythonComponent( # dummy python_func not used in behavior that is being tested python_func=lambda: None, component_spec=original_component_spec, ) with tempfile.TemporaryDirectory() as tempdir: output_path = os.path.join(tempdir, 'component.yaml') compiler.Compiler().compile(test_component, output_path) # test that it can be read back correctly with open(output_path, 'r') as f: contents = f.read() new_component_spec = structures.ComponentSpec.from_yaml_documents( contents) self.assertEqual(original_component_spec, new_component_spec) def test_simple_component_spec_load_from_v2_component_yaml(self): component_yaml_v2 = textwrap.dedent("""\ components: comp-component-1: executorLabel: exec-component-1 inputDefinitions: parameters: input1: parameterType: STRING outputDefinitions: parameters: output1: parameterType: STRING deploymentSpec: executors: exec-component-1: container: command: - sh - -c - 'set -ex echo "$0" > "$1"' - '{{$.inputs.parameters[''input1'']}}' - '{{$.outputs.parameters[''output1''].output_file}}' image: alpine pipelineInfo: name: component-1 root: dag: tasks: component-1: cachingOptions: enableCache: true componentRef: name: comp-component-1 inputs: parameters: input1: componentInputParameter: input1 taskInfo: name: component-1 inputDefinitions: parameters: input1: parameterType: STRING schemaVersion: 2.1.0 sdkVersion: kfp-2.0.0-alpha.2 """) generated_spec = structures.ComponentSpec.from_yaml_documents( component_yaml_v2) expected_spec = structures.ComponentSpec( name='component-1', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'sh', '-c', 'set -ex\necho "$0" > "$1"', "{{$.inputs.parameters['input1']}}", "{{$.outputs.parameters['output1'].output_file}}", ], )), inputs={'input1': structures.InputSpec(type='String')}, outputs={'output1': structures.OutputSpec(type='String')}) self.assertEqual(generated_spec, expected_spec) @parameterized.parameters( { 'yaml': V1_YAML_IF_PLACEHOLDER, 'expected_component': COMPONENT_SPEC_IF_PLACEHOLDER }, { 'yaml': V1_YAML_CONCAT_PLACEHOLDER, 'expected_component': COMPONENT_SPEC_CONCAT_PLACEHOLDER }, { 'yaml': V1_YAML_NESTED_PLACEHOLDER, 'expected_component': COMPONENT_SPEC_NESTED_PLACEHOLDER }, { 'yaml': V1_YAML_EXECUTOR_INPUT_PLACEHOLDER, 'expected_component': COMPONENT_SPEC_EXECUTOR_INPUT_PLACEHOLDER }, ) def test_component_spec_placeholder_load_from_v2_component_yaml( self, yaml, expected_component): generated_spec = structures.ComponentSpec.from_yaml_documents(yaml) self.assertEqual(generated_spec, expected_component) def test_component_spec_load_from_v1_component_yaml(self): component_yaml_v1 = textwrap.dedent("""\ name: Component with 2 inputs and 2 outputs inputs: - {name: Input parameter, type: String} - {name: Input artifact} outputs: - {name: Output 1} - {name: Output 2} implementation: container: image: busybox command: [sh, -c, ' mkdir -p $(dirname "$2") mkdir -p $(dirname "$3") echo "$0" > "$2" cp "$1" "$3" ' ] args: - {inputValue: Input parameter} - {inputPath: Input artifact} - {outputPath: Output 1} - {outputPath: Output 2} """) generated_spec = structures.ComponentSpec.from_yaml_documents( component_yaml_v1) expected_spec = structures.ComponentSpec( name='Component with 2 inputs and 2 outputs', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='busybox', command=[ 'sh', '-c', (' mkdir -p $(dirname "$2") mkdir -p $(dirname "$3") ' 'echo "$0" > "$2" cp "$1" "$3" '), ], args=[ placeholders.InputValuePlaceholder( input_name='input_parameter'), placeholders.InputPathPlaceholder( input_name='input_artifact'), placeholders.OutputPathPlaceholder( output_name='output_1'), placeholders.OutputPathPlaceholder( output_name='output_2'), ], env={}, )), inputs={ 'input_parameter': structures.InputSpec(type='String'), 'input_artifact': structures.InputSpec(type='system.Artifact@0.0.1') }, outputs={ 'output_1': structures.OutputSpec(type='system.Artifact@0.0.1'), 'output_2': structures.OutputSpec(type='system.Artifact@0.0.1'), }) self.assertEqual(generated_spec, expected_spec) class TestContainerSpecImplementation(unittest.TestCase): def test_command_and_args(self): obj = structures.ContainerSpecImplementation( image='image', command=['command'], args=['args']) self.assertEqual(obj.command, ['command']) self.assertEqual(obj.args, ['args']) obj = structures.ContainerSpecImplementation( image='image', command=[], args=[]) self.assertEqual(obj.command, None) self.assertEqual(obj.args, None) def test_env(self): obj = structures.ContainerSpecImplementation( image='image', command=['command'], args=['args'], env={'env': 'env'}) self.assertEqual(obj.env, {'env': 'env'}) obj = structures.ContainerSpecImplementation( image='image', command=[], args=[], env={}) self.assertEqual(obj.env, None) def test_from_container_dict_no_placeholders(self): expected_container_spec = structures.ContainerSpecImplementation( image='python:3.7', command=['sh', '-c', 'dummy'], args=['--executor_input', '{{$}}', '--function_to_execute', 'func'], env={'ENV1': 'val1'}, resources=None) container_dict = { 'args': [ '--executor_input', '{{$}}', '--function_to_execute', 'func' ], 'command': ['sh', '-c', 'dummy'], 'image': 'python:3.7', 'env': { 'ENV1': 'val1' }, } loaded_container_spec = structures.ContainerSpecImplementation.from_container_dict( container_dict) self.assertEqual(expected_container_spec, loaded_container_spec) def test_raise_error_if_access_artifact_by_itself(self): def comp_with_artifact_input(dataset: dsl.Input[dsl.Dataset]): return dsl.ContainerSpec( image='gcr.io/my-image', command=['sh', 'run.sh'], args=[dataset]) def comp_with_artifact_output(dataset_old: dsl.Output[dsl.Dataset], dataset_new: dsl.Output[dsl.Dataset], optional_input: str = 'default'): return dsl.ContainerSpec( image='gcr.io/my-image', command=['sh', 'run.sh'], args=[ dsl.IfPresentPlaceholder( input_name='optional_input', then=[dataset_old], else_=[dataset_new]) ]) self.assertRaisesRegex( ValueError, r'Cannot access artifact by itself in the container definition.', component_factory.create_container_component_from_func, comp_with_artifact_input) self.assertRaisesRegex( ValueError, r'Cannot access artifact by itself in the container definition.', component_factory.create_container_component_from_func, comp_with_artifact_output) class TestComponentSpec(unittest.TestCase): def test_inputs(self): obj = structures.ComponentSpec( name='name', implementation=structures.Implementation(container=None), inputs={}) self.assertEqual(obj.inputs, None) def test_outputs(self): obj = structures.ComponentSpec( name='name', implementation=structures.Implementation(container=None), outputs={}) self.assertEqual(obj.outputs, None) class TestInputSpec(unittest.TestCase): def test_equality(self): self.assertEqual( structures.InputSpec(type='String', default=None), structures.InputSpec(type='String', default=None)) self.assertNotEqual( structures.InputSpec(type='String', default=None), structures.InputSpec(type='String', default='test', optional=True)) self.assertEqual( structures.InputSpec(type='List', default=None), structures.InputSpec(type='typing.List', default=None)) self.assertEqual( structures.InputSpec(type='List', default=None), structures.InputSpec(type='typing.List[int]', default=None)) self.assertEqual( structures.InputSpec(type='List'), structures.InputSpec(type='typing.List[typing.Dict[str, str]]')) def test_optional(self): input_spec = structures.InputSpec( type='String', default=None, optional=True) self.assertEqual(input_spec.default, None) self.assertEqual(input_spec.optional, True) input_spec = structures.InputSpec(type='String') self.assertEqual(input_spec.default, None) self.assertEqual(input_spec.optional, False) def test_from_ir_component_inputs_dict(self): parameter_dict = {'parameterType': 'STRING'} input_spec = structures.InputSpec.from_ir_component_inputs_dict( parameter_dict) self.assertEqual(input_spec.type, 'String') self.assertEqual(input_spec.default, None) parameter_dict = {'parameterType': 'NUMBER_INTEGER'} input_spec = structures.InputSpec.from_ir_component_inputs_dict( parameter_dict) self.assertEqual(input_spec.type, 'Integer') self.assertEqual(input_spec.default, None) parameter_dict = { 'defaultValue': 'default value', 'parameterType': 'STRING' } input_spec = structures.InputSpec.from_ir_component_inputs_dict( parameter_dict) self.assertEqual(input_spec.type, 'String') self.assertEqual(input_spec.default, 'default value') input_spec = structures.InputSpec.from_ir_component_inputs_dict( parameter_dict) self.assertEqual(input_spec.type, 'String') self.assertEqual(input_spec.default, 'default value') artifact_dict = { 'artifactType': { 'schemaTitle': 'system.Artifact', 'schemaVersion': '0.0.1' } } input_spec = structures.InputSpec.from_ir_component_inputs_dict( artifact_dict) self.assertEqual(input_spec.type, 'system.Artifact@0.0.1') def test_assert_optional_must_be_true_when_default_is_not_none(self): with self.assertRaisesRegex(ValueError, r'must be True if `default` is not None'): input_spec = structures.InputSpec(type='String', default='text') class TestOutputSpec(parameterized.TestCase): def test_from_ir_component_outputs_dict(self): parameter_dict = {'parameterType': 'STRING'} output_spec = structures.OutputSpec.from_ir_component_outputs_dict( parameter_dict) self.assertEqual(output_spec.type, 'String') artifact_dict = { 'artifactType': { 'schemaTitle': 'system.Artifact', 'schemaVersion': '0.0.1' } } output_spec = structures.OutputSpec.from_ir_component_outputs_dict( artifact_dict) self.assertEqual(output_spec.type, 'system.Artifact@0.0.1') V1_YAML = textwrap.dedent("""\ implementation: container: args: - if: cond: isPresent: optional_input_1 else: - --arg2 - default then: - --arg1 - {inputUri: optional_input_1} image: alpine inputs: - {name: optional_input_1, optional: true, type: String} name: component_if """) class TestReadInComponent(parameterized.TestCase): def test_read_v1(self): component_spec = structures.ComponentSpec.from_yaml_documents( V1_YAML_IF_PLACEHOLDER) self.assertEqual(component_spec.name, 'component-if') self.assertEqual(component_spec.implementation.container.image, 'alpine') def test_simple_placeholder(self): compiled_yaml = textwrap.dedent(""" components: comp-component1: executorLabel: exec-component1 inputDefinitions: parameters: input1: parameterType: STRING outputDefinitions: artifacts: output1: artifactType: schemaTitle: system.Artifact schemaVersion: 0.0.1 deploymentSpec: executors: exec-component1: container: args: - '{{$.inputs.parameters[''input1'']}}' - '{{$.outputs.artifacts[''output1''].path}}' command: - sh - -c - echo "$0" >> "$1" image: alpine pipelineInfo: name: component1 root: dag: tasks: component1: cachingOptions: enableCache: true componentRef: name: comp-component1 inputs: parameters: input1: componentInputParameter: input1 taskInfo: name: component1 inputDefinitions: parameters: input1: parameterType: STRING schemaVersion: 2.1.0 sdkVersion: kfp-2.0.0-alpha.2""") loaded_component_spec = structures.ComponentSpec.from_yaml_documents( compiled_yaml) component_spec = structures.ComponentSpec( name='component1', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=['sh', '-c', 'echo "$0" >> "$1"'], args=[ "{{$.inputs.parameters['input1']}}", "{{$.outputs.artifacts['output1'].path}}", ], env=None, resources=None), graph=None, importer=None), description=None, inputs={ 'input1': structures.InputSpec(type='String', default=None) }, outputs={ 'output1': structures.OutputSpec(type='system.Artifact@0.0.1') }) self.assertEqual(loaded_component_spec, component_spec) def test_if_placeholder(self): compiled_yaml = textwrap.dedent(""" components: comp-if: executorLabel: exec-if inputDefinitions: parameters: optional_input_1: parameterType: STRING deploymentSpec: executors: exec-if: container: args: - 'input: ' - '{{$.inputs.parameters[''optional_input_1'']}}' command: - sh - -c - echo "$0" "$1" image: alpine pipelineInfo: name: if root: dag: tasks: if: cachingOptions: enableCache: true componentRef: name: comp-if inputs: parameters: optional_input_1: componentInputParameter: optional_input_1 taskInfo: name: if inputDefinitions: parameters: optional_input_1: parameterType: STRING schemaVersion: 2.1.0 sdkVersion: kfp-2.0.0-alpha.2""") loaded_component_spec = structures.ComponentSpec.from_yaml_documents( compiled_yaml) component_spec = structures.ComponentSpec( name='if', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=['sh', '-c', 'echo "$0" "$1"'], args=[ 'input: ', "{{$.inputs.parameters['optional_input_1']}}", ], env=None, resources=None), graph=None, importer=None), description=None, inputs={ 'optional_input_1': structures.InputSpec(type='String', default=None) }, outputs=None) self.assertEqual(loaded_component_spec, component_spec) def test_concat_placeholder(self): compiled_yaml = textwrap.dedent(""" components: comp-concat: executorLabel: exec-concat inputDefinitions: parameters: input1: parameterType: STRING input2: parameterType: STRING deploymentSpec: executors: exec-concat: container: command: - sh - -c - echo "$0" - '{{$.inputs.parameters[''input1'']}}+{{$.inputs.parameters[''input2'']}}' image: alpine pipelineInfo: name: concat root: dag: tasks: concat: cachingOptions: enableCache: true componentRef: name: comp-concat inputs: parameters: input1: componentInputParameter: input1 input2: componentInputParameter: input2 taskInfo: name: concat inputDefinitions: parameters: input1: parameterType: STRING input2: parameterType: STRING schemaVersion: 2.1.0 sdkVersion: kfp-2.0.0-alpha.2""") loaded_component_spec = structures.ComponentSpec.from_yaml_documents( compiled_yaml) component_spec = structures.ComponentSpec( name='concat', implementation=structures.Implementation( container=structures.ContainerSpecImplementation( image='alpine', command=[ 'sh', '-c', 'echo "$0"', "{{$.inputs.parameters['input1']}}+{{$.inputs.parameters['input2']}}", ], args=None, env=None, resources=None), graph=None, importer=None), description=None, inputs={ 'input1': structures.InputSpec(type='String', default=None), 'input2': structures.InputSpec(type='String', default=None) }, outputs=None) self.assertEqual(loaded_component_spec, component_spec) class TestNormalizeTimeString(parameterized.TestCase): @parameterized.parameters([ ('1 hour', '1h'), ('2 hours', '2h'), ('2hours', '2h'), ('2 w', '2w'), ('2d', '2d'), ]) def test(self, unnorm: str, norm: str): self.assertEqual(structures.normalize_time_string(unnorm), norm) def test_multipart_duration_raises(self): with self.assertRaisesRegex(ValueError, 'Invalid duration string:'): structures.convert_duration_to_seconds('1 day 1 hour') def test_non_int_value_raises(self): with self.assertRaisesRegex(ValueError, 'Invalid duration string:'): structures.convert_duration_to_seconds('one hour') class TestConvertDurationToSeconds(parameterized.TestCase): @parameterized.parameters([ ('1 hour', 3600), ('2 hours', 7200), ('2hours', 7200), ('2 w', 1209600), ('2d', 172800), ]) def test(self, duration: str, seconds: int): self.assertEqual( structures.convert_duration_to_seconds(duration), seconds) def test_unsupported_duration_unit(self): with self.assertRaisesRegex(ValueError, 'Unsupported duration unit:'): structures.convert_duration_to_seconds('1 year') class TestRetryPolicy(unittest.TestCase): def test_to_proto(self): retry_policy_struct = structures.RetryPolicy( max_retry_count=10, backoff_duration='1h', backoff_factor=1.5, backoff_max_duration='2 weeks') retry_policy_proto = retry_policy_struct.to_proto() self.assertEqual(retry_policy_proto.max_retry_count, 10) self.assertEqual(retry_policy_proto.backoff_duration.seconds, 3600) self.assertEqual(retry_policy_proto.backoff_factor, 1.5) # tests cap self.assertEqual(retry_policy_proto.backoff_max_duration.seconds, 3600) class TestDeserializeV1ComponentYamlDefaults(unittest.TestCase): def test_True(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Boolean, default: "True" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, True) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.bool_value, True) def test_true(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Boolean, default: "true" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, True) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.bool_value, True) def test_false(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Boolean, default: "false" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, False) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.bool_value, False) def test_False(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Boolean, default: "False" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, False) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.bool_value, False) def test_int(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Integer, default: "1" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, 1) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.number_value, 1.0) def test_float(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: Float, default: "1.0" } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, 1.0) self.assertEqual( comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.number_value, 1.0) def test_struct(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: JsonObject, default: '{"a": 1.0}' } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, {'a': 1.0}) self.assertEqual( dict(comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.struct_value), {'a': 1.0}) def test_array(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: JsonObject, default: '["a", 1.0]' } implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputValue: val } """) comp = components.load_component_from_text(comp_text) self.assertEqual(comp.component_spec.inputs['val'].default, ['a', 1.0]) self.assertEqual( list(comp.pipeline_spec.root.input_definitions.parameters['val'] .default_value.list_value), ['a', 1.0]) def test_artifact_with_dict_type_passes_through(self): comp_text = textwrap.dedent("""\ inputs: - { name: val, type: {Key: Val}} implementation: container: image: alpine command: - sh - -c - | echo $0 - { inputPath: val } """) comp = components.load_component_from_text(comp_text) self.assertFalse(comp.component_spec.inputs['val'].optional) self.assertFalse(comp.pipeline_spec.root.input_definitions .artifacts['val'].is_optional) def test_load_noncanonical_v1_generic_types(self): loaded_comp = components.load_component_from_text( V1_NONCANONICAL_GENERIC_TYPES_COMPONENT_SPEC) inputs = loaded_comp.component_spec.inputs outputs = loaded_comp.component_spec.outputs self.assertEqual(inputs['input1'].type, 'List') self.assertEqual(inputs['input2'].type, 'List') self.assertEqual(inputs['input3'].type, 'Dict') self.assertEqual(inputs['input4'].type, 'Dict') self.assertEqual(outputs['output1'].type, 'List') self.assertEqual(outputs['output2'].type, 'List') self.assertEqual(outputs['output3'].type, 'Dict') self.assertEqual(outputs['output4'].type, 'Dict') class TestLoadDocumentsFromYAML(unittest.TestCase): def test_no_documents(self): with self.assertRaisesRegex( ValueError, r'Expected one or two YAML documents in the IR YAML file\. Got\: 0\.' ): structures.load_documents_from_yaml('') def test_one_document(self): doc1, doc2 = structures.load_documents_from_yaml( textwrap.dedent("""\ key1: value1 """)) self.assertEqual(doc1, {'key1': 'value1'}) self.assertEqual(doc2, {}) def test_two_documents(self): doc1, doc2 = structures.load_documents_from_yaml( textwrap.dedent("""\ key1: value1 --- key2: value2 """)) self.assertEqual(doc1, {'key1': 'value1'}) self.assertEqual(doc2, {'key2': 'value2'}) def test_three_documents(self): with self.assertRaisesRegex( ValueError, r'Expected one or two YAML documents in the IR YAML file\. Got\: 3\.' ): structures.load_documents_from_yaml( textwrap.dedent("""\ key3: value3 --- key3: value3 --- key3: value3 """)) if __name__ == '__main__': unittest.main()