def test_model_modify_node():
node = ast_call("TensorFlowModel(image=my_image)")
modifier = renamed_params.ModelImageURIRenamer()
modifier.modify_node(node)
expected = "TensorFlowModel(image_uri=my_image)"
assert expected == pasta.dump(node)
def test_model_node_should_be_modified_no_distribution():
modifier = renamed_params.ModelImageURIRenamer()
for model, namespaces in MODELS.items():
call = "{}()".format(model)
assert not modifier.node_should_be_modified(ast_call(call))
for namespace in namespaces:
call = "{}.{}()".format(namespace, model)
assert not modifier.node_should_be_modified(ast_call(call))
def test_model_node_should_be_modified():
modifier = renamed_params.ModelImageURIRenamer()
for model, namespaces in MODELS.items():
call = "{}(image='my-image:latest')".format(model)
assert modifier.node_should_be_modified(ast_call(call))
for namespace in namespaces:
call = "{}.{}(image='my-image:latest')".format(namespace, model)
assert modifier.node_should_be_modified(ast_call(call))
def test_model_node_should_be_modified_random_function_call():
modifier = renamed_params.ModelImageURIRenamer()
assert not modifier.node_should_be_modified(ast_call("Session()"))