1 files changed, 2 insertions, 2 deletions
diff --git a/unittests/Tooling/TestVisitor.h b/unittests/Tooling/TestVisitor.h
index f73d2e0c2f..d439d81d89 100644
--- a/unittests/Tooling/TestVisitor.h
+++ b/unittests/Tooling/TestVisitor.h
@@ -56,6 +56,7 @@ protected:
     FindConsumer(TestVisitor *Visitor) : Visitor(Visitor) {}
 
     virtual void HandleTranslationUnit(clang::ASTContext &Context) {
+      Visitor->Context = &Context;
       Visitor->TraverseDecl(Context.getTranslationUnitDecl());
     }
 
@@ -68,8 +69,7 @@ protected:
     TestAction(TestVisitor *Visitor) : Visitor(Visitor) {}
 
     virtual clang::ASTConsumer* CreateASTConsumer(
-        CompilerInstance& compiler, llvm::StringRef dummy) {
-      Visitor->Context = &compiler.getASTContext();
+        CompilerInstance&, llvm::StringRef dummy) {
       /// TestConsumer will be deleted by the framework calling us.
       return new FindConsumer(Visitor);
     }