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))
