def test_import_from_modify_node(import_statements):
modifier = image_uris.ImageURIRetrieveImportFromRenamer()
expected_result = "from sagemaker import image_uris"
for import_statement in import_statements:
node = ast_import(import_statement)
modifier.modify_node(node)
assert expected_result == pasta.dump(node)
def test_import_from_node_should_be_modified_image_uris_input(import_statements):
modifier = image_uris.ImageURIRetrieveImportFromRenamer()
statement = "from sagemaker import get_image_uri"
node = ast_import(statement)
assert modifier.node_should_be_modified(node)
statement = "from sagemaker.amazon_estimator import get_image_uri"
node = ast_import(statement)
assert modifier.node_should_be_modified(node)
statement = "from sagemaker.amazon.amazon_estimator import get_image_uri"
node = ast_import(statement)
assert modifier.node_should_be_modified(node)
def test_import_from_node_should_be_modified_random_import():
modifier = image_uris.ImageURIRetrieveImportFromRenamer()
node = ast_import("from sagemaker.amazon_estimator import registry")
assert not modifier.node_should_be_modified(node)