1 files changed, 18 insertions, 9 deletions
diff --git a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
index c2a1b7f44eb..3daed75d5a7 100644
--- a/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
+++ b/clang-tools-extra/clangd/unittests/TUSchedulerTests.cpp
@@ -7,10 +7,14 @@
 //===----------------------------------------------------------------------===//
 
 #include "Annotations.h"
+#include "ClangdUnit.h"
 #include "Context.h"
+#include "Diagnostics.h"
 #include "Matchers.h"
+#include "Path.h"
 #include "TUScheduler.h"
 #include "TestFS.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
@@ -56,12 +60,17 @@ protected:
   /// in updateWithDiags.
   static std::unique_ptr<ParsingCallbacks> captureDiags() {
     class CaptureDiags : public ParsingCallbacks {
-      void onDiagnostics(PathRef File, std::vector<Diag> Diags) override {
+      void onMainAST(PathRef File, ParsedAST &AST, PublishFn Publish) override {
+        auto Diags = AST.getDiagnostics();
         auto D = Context::current().get(DiagsCallbackKey);
         if (!D)
           return;
-        const_cast<llvm::unique_function<void(PathRef, std::vector<Diag>)> &> (
-            *D)(File, Diags);
+
+        Publish([&]() {
+          const_cast<
+              llvm::unique_function<void(PathRef, std::vector<Diag>)> &> (*D)(
+              File, std::move(Diags));
+        });
       }
     };
     return llvm::make_unique<CaptureDiags>();
@@ -116,8 +125,8 @@ TEST_F(TUSchedulerTests, MissingFiles) {
   S.update(Added, getInputs(Added, "x"), WantDiagnostics::No);
   EXPECT_EQ(S.getContents(Added), "x");
 
-  // Assert each operation for missing file is an error (even if it's available
-  // in VFS).
+  // Assert each operation for missing file is an error (even if it's
+  // available in VFS).
   S.runWithAST("", Missing,
                [&](Expected<InputsAndAST> AST) { EXPECT_ERROR(AST); });
   S.runWithPreamble(
@@ -367,8 +376,8 @@ TEST_F(TUSchedulerTests, ManyUpdates) {
     StringRef AllContents[] = {Contents1, Contents2, Contents3};
     const int AllContentsSize = 3;
 
-    // Scheduler may run tasks asynchronously, but should propagate the context.
-    // We stash a nonce in the context, and verify it in the task.
+    // Scheduler may run tasks asynchronously, but should propagate the
+    // context. We stash a nonce in the context, and verify it in the task.
     static Key<int> NonceKey;
     int Nonce = 0;
 
@@ -465,8 +474,8 @@ TEST_F(TUSchedulerTests, EvictedAST) {
   ASSERT_TRUE(S.blockUntilIdle(timeoutSeconds(10)));
   ASSERT_EQ(BuiltASTCounter.load(), 1);
 
-  // Build two more files. Since we can retain only 2 ASTs, these should be the
-  // ones we see in the cache later.
+  // Build two more files. Since we can retain only 2 ASTs, these should be
+  // the ones we see in the cache later.
   updateWithCallback(S, Bar, SourceContents, WantDiagnostics::Yes,
                      [&BuiltASTCounter]() { ++BuiltASTCounter; });
   updateWithCallback(S, Baz, SourceContents, WantDiagnostics::Yes,