def test_model_artifact(self):
model = ontology_artifacts.Model()
model.framework = 'TFX'
model.metadata["framework_version"] = '1.12'
model.metadata['custom_field'] = {
'custom_field_1': 1,
'custom_field_2': [1, 3, 4]
}
with open(
os.path.join(os.path.dirname(__file__), 'test_data',
'expected_model_artifact.json')) as json_file:
expected_model_json = json.load(json_file)
self.assertEqual(expected_model_json,
json.loads(model.serialize()))
def testSerialization(self):
my_model = ontology_artifacts.Model()
my_model.set_float_custom_property('float1', 1.1)
my_model.set_int_custom_property('int1', 1)
my_model.set_string_custom_property('string1', 'testString')
self.assertEqual(_EXPECTED_SERIALIZATION, my_model.serialize())
rehydrated_model = artifact.Artifact.deserialize(
_EXPECTED_SERIALIZATION)
self.assertEqual(ontology_artifacts.Model, type(rehydrated_model))
self.assertEqual(
'testString',
rehydrated_model.get_string_custom_property('string1'))
self.assertEqual(1, rehydrated_model.get_int_custom_property('int1'))
self.assertEqual(1.1,
rehydrated_model.get_float_custom_property('float1'))
def testGetExecutorOutput(self):
model = ontology_artifacts.Model()
model.name = 'test-artifact'
model.uri = 'gs://root/execution/output'
model.metadata['test_property'] = 'test value'
executor_output = entrypoint_utils.get_executor_output(
output_artifacts={'output': model},
output_params={
'int_output': 42,
'string_output': 'hello world!',
'float_output': 12.12
})
# Renormalize the JSON proto read from testdata. Otherwise there'll be
# mismatch in the way treating int value.
expected_output = pipeline_spec_pb2.ExecutorOutput()
expected_output = json_format.Parse(
text=_get_text_from_testdata('executor_output.json'),
message=expected_output)
self.assertDictEqual(json_format.MessageToDict(expected_output),
json_format.MessageToDict(executor_output))