1 files changed, 33 insertions, 6 deletions

diff --git a/unittests/CrossTU/CrossTranslationUnitTest.cpp b/unittests/CrossTU/CrossTranslationUnitTest.cpp
index b3e7243ce1..b4f22d5cb0 100644
--- a/unittests/CrossTU/CrossTranslationUnitTest.cpp
+++ b/unittests/CrossTU/CrossTranslationUnitTest.cpp
@@ -28,13 +28,18 @@ public:
       : CTU(CI), Success(Success) {}
 
   void HandleTranslationUnit(ASTContext &Ctx) {
+    auto FindFInTU = [](const TranslationUnitDecl *TU) {
+      const FunctionDecl *FD = nullptr;
+      for (const Decl *D : TU->decls()) {
+        FD = dyn_cast<FunctionDecl>(D);
+        if (FD && FD->getName() == "f")
+          break;
+      }
+      return FD;
+    };
+
     const TranslationUnitDecl *TU = Ctx.getTranslationUnitDecl();
-    const FunctionDecl *FD = nullptr;
-    for (const Decl *D : TU->decls()) {
-      FD = dyn_cast<FunctionDecl>(D);
-      if (FD && FD->getName() == "f")
-        break;
-    }
+    const FunctionDecl *FD = FindFInTU(TU);
     assert(FD && FD->getName() == "f");
     bool OrigFDHasBody = FD->hasBody();
 
@@ -78,6 +83,28 @@ public:
     if (NewFDorError) {
       const FunctionDecl *NewFD = *NewFDorError;
       *Success = NewFD && NewFD->hasBody() && !OrigFDHasBody;
+
+      if (NewFD) {
+        // Check GetImportedFromSourceLocation.
+        llvm::Optional<std::pair<SourceLocation, ASTUnit *>> SLocResult =
+            CTU.getImportedFromSourceLocation(NewFD->getLocation());
+        EXPECT_TRUE(SLocResult);
+        if (SLocResult) {
+          SourceLocation OrigSLoc = (*SLocResult).first;
+          ASTUnit *OrigUnit = (*SLocResult).second;
+          // OrigUnit is created internally by CTU (is not the
+          // ASTWithDefinition).
+          TranslationUnitDecl *OrigTU =
+              OrigUnit->getASTContext().getTranslationUnitDecl();
+          const FunctionDecl *FDWithDefinition = FindFInTU(OrigTU);
+          EXPECT_TRUE(FDWithDefinition);
+          if (FDWithDefinition) {
+            EXPECT_EQ(FDWithDefinition->getName(), "f");
+            EXPECT_TRUE(FDWithDefinition->isThisDeclarationADefinition());
+            EXPECT_EQ(OrigSLoc, FDWithDefinition->getLocation());
+          }
+        }
+      }
     }
   }