190 lines
6.6 KiB
Python
190 lines
6.6 KiB
Python
# Copyright 2020 Google LLC
|
|
#
|
|
# 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.artifact module."""
|
|
import unittest
|
|
import textwrap
|
|
|
|
from kfp.dsl import artifact
|
|
|
|
|
|
class _MyArtifact(artifact.Artifact):
|
|
TYPE_NAME = 'MyTypeName'
|
|
PROPERTIES = {
|
|
'int1': artifact.Property(
|
|
type=artifact.PropertyType.INT,
|
|
description='An integer-typed property'),
|
|
'int2': artifact.Property(type=artifact.PropertyType.INT),
|
|
'float1': artifact.Property(
|
|
type=artifact.PropertyType.DOUBLE,
|
|
description='A float-typed property'),
|
|
'float2': artifact.Property(type=artifact.PropertyType.DOUBLE),
|
|
'string1': artifact.Property(
|
|
type=artifact.PropertyType.STRING,
|
|
description='A string-typed property'),
|
|
'string2': artifact.Property(type=artifact.PropertyType.STRING),
|
|
}
|
|
|
|
|
|
_SERIALIZED_INSTANCE = """\
|
|
{
|
|
"properties": {
|
|
"float1": {
|
|
"doubleValue": 1.11
|
|
},
|
|
"int1": {
|
|
"intValue": "1"
|
|
},
|
|
"string1": {
|
|
"stringValue": "111"
|
|
}
|
|
},
|
|
"type": {
|
|
"instanceSchema": "properties:\\n float1:\\n description: A float-typed property\\n type: double\\n float2:\\n description:\\n type: double\\n int1:\\n description: An integer-typed property\\n type: int\\n int2:\\n description:\\n type: int\\n string1:\\n description: A string-typed property\\n type: string\\n string2:\\n description:\\n type: string\\ntitle: kfp.MyTypeName\\ntype: object\\n"
|
|
}
|
|
}"""
|
|
|
|
|
|
class ArtifactTest(unittest.TestCase):
|
|
|
|
def testArtifact(self):
|
|
instance = _MyArtifact()
|
|
|
|
# Test property getters.
|
|
self.assertEqual('', instance.uri)
|
|
self.assertEqual('', instance.name)
|
|
|
|
# Default property does not have span or split_names.
|
|
with self.assertRaisesRegex(AttributeError, "has no property 'span'"):
|
|
_ = instance.span
|
|
with self.assertRaisesRegex(AttributeError,
|
|
"has no property 'split_names'"):
|
|
_ = instance.split_names
|
|
|
|
# Test property setters.
|
|
instance.uri = '/tmp/uri2'
|
|
self.assertEqual('/tmp/uri2', instance.uri)
|
|
|
|
instance.name = '1'
|
|
self.assertEqual('1', instance.name)
|
|
|
|
# Testing artifact does not have span.
|
|
with self.assertRaisesRegex(AttributeError, "unknown property 'span'"):
|
|
instance.span = 20190101
|
|
# Testing artifact does not have span.
|
|
with self.assertRaisesRegex(AttributeError,
|
|
"unknown property 'split_names'"):
|
|
instance.split_names = ''
|
|
|
|
instance.set_int_custom_property('int_key', 20)
|
|
self.assertEqual(
|
|
20, instance.runtime_artifact.custom_properties['int_key'].int_value)
|
|
|
|
instance.set_string_custom_property('string_key', 'string_value')
|
|
self.assertEqual(
|
|
'string_value',
|
|
instance.runtime_artifact.custom_properties['string_key'].string_value)
|
|
|
|
self.assertEqual(textwrap.dedent("""\
|
|
Artifact(artifact: name: "1"
|
|
type {
|
|
instance_schema: "properties:\\n float1:\\n description: A float-typed property\\n type: double\\n float2:\\n description:\\n type: double\\n int1:\\n description: An integer-typed property\\n type: int\\n int2:\\n description:\\n type: int\\n string1:\\n description: A string-typed property\\n type: string\\n string2:\\n description:\\n type: string\\ntitle: kfp.MyTypeName\\ntype: object\\n"
|
|
}
|
|
uri: "/tmp/uri2"
|
|
custom_properties {
|
|
key: "int_key"
|
|
value {
|
|
int_value: 20
|
|
}
|
|
}
|
|
custom_properties {
|
|
key: "string_key"
|
|
value {
|
|
string_value: "string_value"
|
|
}
|
|
}
|
|
, type_schema: properties:
|
|
float1:
|
|
description: A float-typed property
|
|
type: double
|
|
float2:
|
|
description:
|
|
type: double
|
|
int1:
|
|
description: An integer-typed property
|
|
type: int
|
|
int2:
|
|
description:
|
|
type: int
|
|
string1:
|
|
description: A string-typed property
|
|
type: string
|
|
string2:
|
|
description:
|
|
type: string
|
|
title: kfp.MyTypeName
|
|
type: object
|
|
)"""), str(instance))
|
|
|
|
def testArtifactProperties(self):
|
|
my_artifact = _MyArtifact()
|
|
|
|
self.assertEqual(0, my_artifact.int1)
|
|
self.assertEqual(0, my_artifact.int2)
|
|
my_artifact.int1 = 111
|
|
my_artifact.int2 = 222
|
|
self.assertEqual('', my_artifact.string1)
|
|
self.assertEqual('', my_artifact.string2)
|
|
my_artifact.string1 = '111'
|
|
my_artifact.string2 = '222'
|
|
self.assertEqual(0.0, my_artifact.float1)
|
|
self.assertEqual(0.0, my_artifact.float2)
|
|
my_artifact.float1 = 1.11
|
|
my_artifact.float2 = 2.22
|
|
self.assertEqual(my_artifact.int1, 111)
|
|
self.assertEqual(my_artifact.int2, 222)
|
|
self.assertEqual(my_artifact.string1, '111')
|
|
self.assertEqual(my_artifact.string2, '222')
|
|
self.assertEqual(1.11, my_artifact.float1)
|
|
self.assertEqual(2.22, my_artifact.float2)
|
|
self.assertEqual(my_artifact.get_string_custom_property('invalid'), '')
|
|
self.assertEqual(my_artifact.get_int_custom_property('invalid'), 0)
|
|
self.assertNotIn('invalid', my_artifact._artifact.custom_properties)
|
|
|
|
with self.assertRaisesRegex(
|
|
AttributeError, "Cannot set unknown property 'invalid' on artifact"):
|
|
my_artifact.invalid = 1
|
|
|
|
with self.assertRaisesRegex(
|
|
AttributeError, "Cannot set unknown property 'invalid' on artifact"):
|
|
my_artifact.invalid = 'x'
|
|
|
|
with self.assertRaisesRegex(AttributeError,
|
|
"\D+ artifact has no property 'invalid'"):
|
|
_ = my_artifact.invalid
|
|
|
|
def testSerialize(self):
|
|
instance = _MyArtifact()
|
|
instance.int1 = 1
|
|
instance.string1 = '111'
|
|
instance.float1 = 1.11
|
|
|
|
self.assertEqual(_SERIALIZED_INSTANCE, instance.serialize())
|
|
|
|
def testDeserialize(self):
|
|
instance = artifact.Artifact.deserialize(_SERIALIZED_INSTANCE)
|
|
self.assertEqual(1, instance.int1)
|
|
self.assertEqual('111', instance.string1)
|
|
self.assertEqual(1.11, instance.float1)
|
|
self.assertEqual('kfp.MyTypeName', instance.type_name)
|