def testEnableCache(self):
model = standard_artifacts.Model()
serving_spec = infra_validator_pb2.ServingSpec()
validation_spec = infra_validator_pb2.ValidationSpec()
infra_validator_1 = component.InfraValidator(
model=channel_utils.as_channel([model]),
serving_spec=serving_spec,
validation_spec=validation_spec)
self.assertEqual(None, infra_validator_1.enable_cache)
infra_validator_2 = component.InfraValidator(
model=channel_utils.as_channel([model]),
serving_spec=serving_spec,
validation_spec=validation_spec,
enable_cache=True)
self.assertEqual(True, infra_validator_2.enable_cache)
def testConstruct(self):
model = standard_artifacts.Model()
serving_spec = infra_validator_pb2.ServingSpec()
infra_validator = component.InfraValidator(
model=channel_utils.as_channel([model]), serving_spec=serving_spec)
# Check channels have been created with proper type.
self.assertEqual(standard_artifacts.Model,
infra_validator.inputs['model'].type)
self.assertEqual(standard_artifacts.InfraBlessing,
infra_validator.outputs['blessing'].type)
# Check exec_properties have been populated.
self.assertEqual(
'{}', # Empty dictionary
infra_validator.exec_properties['serving_spec'])
def testConstruct(self):
model = standard_artifacts.Model()
serving_spec = infra_validator_pb2.ServingSpec()
validation_spec = infra_validator_pb2.ValidationSpec()
infra_validator = component.InfraValidator(
model=channel_utils.as_channel([model]),
serving_spec=serving_spec,
validation_spec=validation_spec)
# Check channels have been created with proper type.
self.assertEqual(
standard_artifacts.Model,
infra_validator.inputs[standard_component_specs.MODEL_KEY].type)
self.assertEqual(
standard_artifacts.InfraBlessing,
infra_validator.outputs[standard_component_specs.BLESSING_KEY].type)
# Check exec_properties have been populated.
self.assertIn(standard_component_specs.SERVING_SPEC_KEY,
infra_validator.exec_properties)
self.assertIn(standard_component_specs.VALIDATION_SPEC_KEY,
infra_validator.exec_properties)