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