[clangd][Tweaks] Improve test infra for testing tweaks with an index

clang-tools-extra/clangd/unittests/tweaks/DefineOutlineTests.cpp

@@ -15,6 +15,8 @@ namespace clang {
 namespace clangd {
 namespace {
 
+using ::testing::UnorderedElementsAre;
+
 TWEAK_TEST(DefineOutline);
 
 TEST_F(DefineOutlineTest, TriggersOnFunctionDecl) {
@@ -747,6 +749,43 @@ TEST_F(DefineOutlineTest, FailsMacroSpecifier) {
   }
 }
 
+TWEAK_WORKSPACE_TEST(DefineOutline);
+
+// Test that DefineOutline's use of getCorrespondingHeaderOrSource()
+// to find the source file corresponding to a header file in which the
+// tweak is invoked is working as intended.
+TEST_F(DefineOutlineWorkspaceTest, FindsCorrespondingSource) {
+  llvm::Annotations HeaderBefore(R"cpp(
+class A {
+  void bar();
+  void f^oo(){}
+};
+)cpp");
+  std::string SourceBefore(R"cpp(
+#include "a.hpp"
+void A::bar(){}
+)cpp");
+  std::string HeaderAfter = R"cpp(
+class A {
+  void bar();
+  void foo();
+};
+)cpp";
+  std::string SourceAfter = R"cpp(
+#include "a.hpp"
+void A::bar(){}
+void A::foo(){}
+)cpp";
+  Workspace.addSource("a.hpp", HeaderBefore.code());
+  Workspace.addMainFile("a.cpp", SourceBefore);
+  auto Result = apply("a.hpp", {HeaderBefore.point(), HeaderBefore.point()});
+  EXPECT_THAT(Result,
+              AllOf(withStatus("success"),
+                    editedFiles(UnorderedElementsAre(
+                        FileWithContents(testPath("a.hpp"), HeaderAfter),
+                        FileWithContents(testPath("a.cpp"), SourceAfter)))));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/tweaks/TweakTesting.cpp

@@ -162,5 +162,37 @@ std::string TweakTest::decorate(llvm::StringRef Code,
       .str();
 }
 
+// TODO: Reuse more code between TweakTest::apply() and
+// TweakWorkspaceTest::apply().
+TweakResult
+TweakWorkspaceTest::apply(StringRef InvocationFile,
+                          llvm::Annotations::Range InvocationRange) {
+  auto AST = Workspace.openFile(InvocationFile);
+  if (!AST) {
+    ADD_FAILURE() << "No file '" << InvocationFile << "' in workspace";
+    return TweakResult{"failed to setup"};
+  }
+
+  auto Index = Workspace.index();
+  auto Result = applyTweak(
+      *AST, InvocationRange, TweakID, Index.get(),
+      &AST->getSourceManager().getFileManager().getVirtualFileSystem());
+  if (!Result)
+    return TweakResult{"unavailable"};
+  if (!*Result)
+    return TweakResult{"fail: " + llvm::toString(Result->takeError())};
+  const auto &Effect = **Result;
+  if ((*Result)->ShowMessage)
+    return TweakResult{"message:\n" + *Effect.ShowMessage};
+
+  TweakResult Retval{"success"};
+  for (auto &It : Effect.ApplyEdits) {
+    auto NewText = It.second.apply();
+    if (!NewText)
+      return TweakResult{"bad edits: " + llvm::toString(NewText.takeError())};
+    Retval.EditedFiles.insert_or_assign(It.first(), *NewText);
+  }
+  return Retval;
+}
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/unittests/tweaks/TweakTesting.h

@@ -10,6 +10,7 @@
 #define LLVM_CLANG_TOOLS_EXTRA_CLANGD_UNITTESTS_TWEAKS_TWEAKTESTING_H
 
 #include "ParsedAST.h"
+#include "TestWorkspace.h"
 #include "index/Index.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
@@ -123,6 +124,73 @@ MATCHER_P2(FileWithContents, FileName, Contents, "") {
 #define EXPECT_AVAILABLE(MarkedCode) EXPECT_AVAILABLE_(MarkedCode, true)
 #define EXPECT_UNAVAILABLE(MarkedCode) EXPECT_AVAILABLE_(MarkedCode, false)
 
+// A helper class to represent the return value of TweakWorkspaceTest::apply().
+struct TweakResult {
+  // A string representation the status of the operation.
+  // For failure cases, this is the same as the return value of
+  // TweakTest::apply() (see the comment above that for details).
+  // For success cases, this is "success".
+  std::string Status;
+  // The contents of all files changed by the tweak, including
+  // the file in which it was invoked. Keys are absolute paths.
+  llvm::StringMap<std::string> EditedFiles = {};
+};
+
+// GTest matchers to allow more easily writing assertions about the
+// expected value of a TweakResult.
+MATCHER_P(withStatus, S, "") { return arg.Status == S; }
+template <class EditedFilesMatcher>
+::testing::Matcher<TweakResult> editedFiles(EditedFilesMatcher M) {
+  return ::testing::Field(&TweakResult::EditedFiles, M);
+}
+
+// Used for formatting TweakResult objects in assertion failure messages,
+// so it's easier to understand what didn't match.
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &Stream,
+                                     const TweakResult &Result) {
+  Stream << "{ status: " << Result.Status << ", editedFiles: [";
+  for (const auto &F : Result.EditedFiles) {
+    Stream << F.first() << ":\n";
+    Stream << F.second;
+  }
+  return Stream << "] }";
+}
+
+// A version of TweakTest that makes it easier to create test cases that
+// involve multiple files which are indexed.
+// Usage:
+//   - Call `Workspace.addMainFile(filename, contents)` to add
+//     source files which are indexer entry points (e.g. would show
+//     up in `compile_commands.json`).
+//   - Call `Workspace.addSource(filename, contents)` to add other
+//     source files (e.g. header files).
+//   - Call `apply(filename, range)` to invoke the tweak on the
+//     indicated file with the given range selected. Can be called
+//     multiple times for the same set of added files.
+// The implementation takes care of building an index reflecting
+// all added source files, and making it available to the tweak.
+// Unlike TweakTest, this does not have a notion of a `CodeContext`
+// (i.e. the contents of all added files are interpreted as being
+// in a File context).
+class TweakWorkspaceTest : public ::testing::Test {
+  const char *TweakID;
+
+public:
+  TweakWorkspaceTest(const char *TweakID) : TweakID(TweakID) {}
+
+  TweakResult apply(StringRef InvocationFile,
+                    llvm::Annotations::Range InvocationRange);
+
+protected:
+  TestWorkspace Workspace;
+};
+
+#define TWEAK_WORKSPACE_TEST(TweakID)                                          \
+  class TweakID##WorkspaceTest : public ::clang::clangd::TweakWorkspaceTest {  \
+  protected:                                                                   \
+    TweakID##WorkspaceTest() : TweakWorkspaceTest(#TweakID) {}                 \
+  }
+
 } // namespace clangd
 } // namespace clang