1 files changed, 37 insertions, 7 deletions
diff --git a/clang/unittests/Format/SortIncludesTest.cpp b/clang/unittests/Format/SortIncludesTest.cpp
index cbbc2baad97..8908eb92a56 100644
--- a/clang/unittests/Format/SortIncludesTest.cpp
+++ b/clang/unittests/Format/SortIncludesTest.cpp
@@ -20,13 +20,15 @@ namespace {
 
 class SortIncludesTest : public ::testing::Test {
 protected:
-  std::string sort(llvm::StringRef Code) {
+  std::string sort(llvm::StringRef Code, StringRef FileName = "input.cpp") {
     std::vector<tooling::Range> Ranges(1, tooling::Range(0, Code.size()));
-    std::string Sorted = applyAllReplacements(
-        Code, sortIncludes(getLLVMStyle(), Code, Ranges, "input.cpp"));
-    return applyAllReplacements(
-        Sorted, reformat(getLLVMStyle(), Sorted, Ranges, "input.cpp"));
+    std::string Sorted =
+        applyAllReplacements(Code, sortIncludes(Style, Code, Ranges, FileName));
+    return applyAllReplacements(Sorted,
+                                reformat(Style, Sorted, Ranges, FileName));
   }
+
+  FormatStyle Style = getLLVMStyle();
 };
 
 TEST_F(SortIncludesTest, BasicSorting) {
@@ -76,13 +78,23 @@ TEST_F(SortIncludesTest, SortsLocallyInEachBlock) {
             "#include \"c.h\"\n"
             "\n"
             "#include \"b.h\"\n",
-            sort("#include \"c.h\"\n"
-                 "#include \"a.h\"\n"
+            sort("#include \"a.h\"\n"
+                 "#include \"c.h\"\n"
                  "\n"
                  "#include \"b.h\"\n"));
 }
 
 TEST_F(SortIncludesTest, HandlesAngledIncludesAsSeparateBlocks) {
+  EXPECT_EQ("#include \"a.h\"\n"
+            "#include \"c.h\"\n"
+            "#include <b.h>\n"
+            "#include <d.h>\n",
+            sort("#include <d.h>\n"
+                 "#include <b.h>\n"
+                 "#include \"c.h\"\n"
+                 "#include \"a.h\"\n"));
+
+  Style = getGoogleStyle(FormatStyle::LK_Cpp);
   EXPECT_EQ("#include <b.h>\n"
             "#include <d.h>\n"
             "#include \"a.h\"\n"
@@ -103,6 +115,24 @@ TEST_F(SortIncludesTest, HandlesMultilineIncludes) {
                  "#include \"b.h\"\n"));
 }
 
+TEST_F(SortIncludesTest, LeavesMainHeaderFirst) {
+  EXPECT_EQ("#include \"llvm/a.h\"\n"
+            "#include \"b.h\"\n"
+            "#include \"c.h\"\n",
+            sort("#include \"llvm/a.h\"\n"
+                 "#include \"c.h\"\n"
+                 "#include \"b.h\"\n"));
+
+  // Don't do this in headers.
+  EXPECT_EQ("#include \"b.h\"\n"
+            "#include \"c.h\"\n"
+            "#include \"llvm/a.h\"\n",
+            sort("#include \"llvm/a.h\"\n"
+                 "#include \"c.h\"\n"
+                 "#include \"b.h\"\n",
+                 "some_header.h"));
+}
+
 } // end namespace
 } // end namespace format
 } // end namespace clang