Change wrapper logic for testing

scripts/flaskext_migrate.py

@@ -83,10 +83,9 @@ def fix_standard_imports(red):
     return red
 
 
-def fix(input_file):
+def fix(ast):
-    ast = read_source(input_file)
+    """Wrapper which allows for testing when not running from shell"""
-    new_ast = fix_imports(ast)
+    return fix_imports(ast).dumps()
-    write_source(new_ast, input_file)
 
 if __name__ == "__main__":
     input_file = sys.argv[1]