diff --git a/tests/check_ast_test.py b/tests/check_ast_test.py
index 64916ba..c16f5fc 100644
--- a/tests/check_ast_test.py
+++ b/tests/check_ast_test.py
@@ -1,15 +1,15 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
 
-from pre_commit_hooks.check_ast import check_ast
+from pre_commit_hooks.check_ast import main
 from testing.util import get_resource_path
 
 
 def test_failing_file():
-    ret = check_ast([get_resource_path('cannot_parse_ast.notpy')])
+    ret = main([get_resource_path('cannot_parse_ast.notpy')])
     assert ret == 1
 
 
 def test_passing_file():
-    ret = check_ast([__file__])
+    ret = main([__file__])
     assert ret == 0