def test_import_check_and_modify_node_random_import():
modifier = tfs.TensorFlowServingImportRenamer()
import_statement = "import random"
node = ast_import(import_statement)
modifier.check_and_modify_node(node)
assert import_statement == pasta.dump(node)
def test_import_from_node_should_be_modified_tfs_module():
import_statements = (
"from sagemaker.tensorflow.serving import Model, Predictor",
"from sagemaker.tensorflow.serving import Predictor",
"from sagemaker.tensorflow.serving import Model as tfsModel",
)
modifier = tfs.TensorFlowServingImportFromRenamer()
for import_from in import_statements:
node = ast_import(import_from)
assert modifier.node_should_be_modified(node) is True
