def testNamedRegistration(self):
@registry.register_model("model2")
class MyModel1(vf_model.VfModel):
pass
model = registry.model("model2")
self.assertTrue(model is MyModel1)
def testNonT2TModelRegistration(self):
@registry.register_model
def model_fn():
pass
model = registry.model("model_fn")
self.assertTrue(model is model_fn)
def testT2TModelRegistration(self):
@registry.register_model
class MyModel1(vf_model.VfModel):
pass
model = registry.model("my_model1")
self.assertTrue(model is MyModel1)
def _model(name):
return registry.model(name)
def testUnknownModel(self):
with self.assertRaisesRegexp(KeyError, "never registered"):
registry.model("not_registered")