1 files changed, 8 insertions, 27 deletions

diff --git a/clang-tools-extra/unittests/clangd/ClangdUnitTests.cpp b/clang-tools-extra/unittests/clangd/ClangdUnitTests.cpp
index ce9859de2d2..6843c9e06c8 100644
--- a/clang-tools-extra/unittests/clangd/ClangdUnitTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdUnitTests.cpp
@@ -10,10 +10,7 @@
 #include "Annotations.h"
 #include "ClangdUnit.h"
 #include "SourceCode.h"
-#include "TestFS.h"
-#include "clang/Frontend/CompilerInvocation.h"
-#include "clang/Frontend/PCHContainerOperations.h"
-#include "clang/Frontend/Utils.h"
+#include "TestTU.h"
 #include "llvm/Support/ScopedPrinter.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -36,24 +33,6 @@ testing::Matcher<const Diag &> WithNote(testing::Matcher<Note> NoteMatcher) {
   return Field(&Diag::Notes, ElementsAre(NoteMatcher));
 }
 
-// FIXME: this is duplicated with FileIndexTests. Share it.
-ParsedAST build(StringRef Code, std::vector<const char *> Flags = {}) {
-  std::vector<const char *> Cmd = {"clang", "main.cpp"};
-  Cmd.insert(Cmd.begin() + 1, Flags.begin(), Flags.end());
-  auto CI = createInvocationFromCommandLine(Cmd);
-  auto Buf = MemoryBuffer::getMemBuffer(Code);
-  auto AST = ParsedAST::Build(std::move(CI), nullptr, std::move(Buf),
-                              std::make_shared<PCHContainerOperations>(),
-                              vfs::getRealFileSystem());
-  assert(AST.hasValue());
-  return std::move(*AST);
-}
-
-std::vector<Diag> buildDiags(llvm::StringRef Code,
-                             std::vector<const char *> Flags = {}) {
-  return build(Code, std::move(Flags)).getDiagnostics();
-}
-
 MATCHER_P2(Diag, Range, Message,
            "Diag at " + llvm::to_string(Range) + " = [" + Message + "]") {
   return arg.Range == Range && arg.Message == Message;
@@ -105,7 +84,7 @@ o]]();
     }
   )cpp");
   EXPECT_THAT(
-      buildDiags(Test.code()),
+      TestTU::withCode(Test.code()).build().getDiagnostics(),
       ElementsAre(
           // This range spans lines.
           AllOf(Diag(Test.range("typo"),
@@ -123,13 +102,15 @@ o]]();
 
 TEST(DiagnosticsTest, FlagsMatter) {
   Annotations Test("[[void]] main() {}");
-  EXPECT_THAT(buildDiags(Test.code()),
+  auto TU = TestTU::withCode(Test.code());
+  EXPECT_THAT(TU.build().getDiagnostics(),
               ElementsAre(AllOf(Diag(Test.range(), "'main' must return 'int'"),
                                 WithFix(Fix(Test.range(), "int",
                                             "change 'void' to 'int'")))));
   // Same code built as C gets different diagnostics.
+  TU.Filename = "Plain.c";
   EXPECT_THAT(
-      buildDiags(Test.code(), {"-x", "c"}),
+      TU.build().getDiagnostics(),
       ElementsAre(AllOf(
           Diag(Test.range(), "return type of 'main' is not 'int'"),
           WithFix(Fix(Test.range(), "int", "change return type to 'int'")))));
@@ -150,7 +131,7 @@ TEST(DiagnosticsTest, Preprocessor) {
     #endif
     )cpp");
   EXPECT_THAT(
-      buildDiags(Test.code()),
+      TestTU::withCode(Test.code()).build().getDiagnostics(),
       ElementsAre(Diag(Test.range(), "use of undeclared identifier 'b'")));
 }
 
@@ -229,7 +210,7 @@ TEST(ClangdUnitTest, GetBeginningOfIdentifier) {
            "int ^λλ^λ();", // UTF-8 handled properly when backing up
        }) {
     Annotations TestCase(Text);
-    auto AST = build(TestCase.code());
+    auto AST = TestTU::withCode(TestCase.code()).build();
     const auto &SourceMgr = AST.getASTContext().getSourceManager();
     SourceLocation Actual = getBeginningOfIdentifier(
         AST, TestCase.points().back(), SourceMgr.getMainFileID());