def __init__(self): examples_channel = channel_utils.as_channel( [simple_artifacts.Dataset()]) external_data_channel = channel_utils.as_channel( [simple_artifacts.File()]) super(ProducerComponent, self).__init__( _ProducerComponentSpec(examples=examples_channel, external_data=external_data_channel))
def testArtifactSchemaMapping(self): # Test first party standard artifact. example_artifact = standard_artifacts.Examples() example_schema = compiler_utils.get_artifact_schema(example_artifact) expected_example_schema = fileio.open( os.path.join(self._schema_base_dir, 'Examples.yaml'), 'rb').read() self.assertEqual(expected_example_schema, example_schema) # Test Kubeflow simple artifact. file_artifact = simple_artifacts.File() file_schema = compiler_utils.get_artifact_schema(file_artifact) expected_file_schema = fileio.open( os.path.join(self._schema_base_dir, 'File.yaml'), 'rb').read() self.assertEqual(expected_file_schema, file_schema) # Test custom artifact type. my_artifact = _MyArtifact() my_artifact_schema = compiler_utils.get_artifact_schema(my_artifact) self.assertDictEqual(yaml.safe_load(my_artifact_schema), yaml.safe_load(_EXPECTED_MY_ARTIFACT_SCHEMA))