def test_add(self, x, y):
    # Test a function with multiple arguments.
    return x + y


TEST_MODELS = {
    "simple-model": SimpleModel,
    "data-structure-model": DataStructureModel
}


def get_model(name):
  if name not in TEST_MODELS:
    raise ValueError("Model name '{}' not in TEST_MODELS")
  return TEST_MODELS[name]()


def main(unused_argv):

  model = get_model(MODEL_NAME.value)