From 1398e40df157101c5e8308dda45ff6d7be348c65 Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Sat, 10 Dec 2022 16:43:32 +0800
Subject: [PATCH] Test `black` compatibility for `.pyi` type stub files

---
 tests/unit/profiles/test_black.py | 86 +++++++++++++++++++++++++++++--
 1 file changed, 82 insertions(+), 4 deletions(-)

diff --git a/tests/unit/profiles/test_black.py b/tests/unit/profiles/test_black.py
index 05444b8cc..a4bb2d52a 100644
--- a/tests/unit/profiles/test_black.py
+++ b/tests/unit/profiles/test_black.py
@@ -19,20 +19,25 @@ def black_format(code: str, is_pyi: bool = False, line_length: int = 88) -> str:
         return code
 
 
-def black_test(code: str, expected_output: str = ""):
+def black_test(code: str, expected_output: str = "", *, is_pyi: bool = False, **config_kwargs):
     """Tests that the given code:
     - Behaves the same when formatted multiple times with isort.
     - Agrees with black formatting.
     - Matches the desired output or itself if none is provided.
     """
     expected_output = expected_output or code
+    config_kwargs = {
+        "extension": "pyi" if is_pyi else None,
+        "profile": "black",
+        **config_kwargs,
+    }
 
     # output should stay consistent over multiple runs
-    output = isort.code(code, profile="black")
-    assert output == isort.code(code, profile="black")
+    output = isort.code(code, **config_kwargs)
+    assert output == isort.code(code, **config_kwargs)
 
     # output should agree with black
-    black_output = black_format(output)
+    black_output = black_format(output, is_pyi=is_pyi)
     assert output == black_output
 
     # output should match expected output
@@ -369,3 +374,76 @@ def test_black_snippet_three():
 DEFAULT_LINE_LENGTH = 88
 """,
     )
+
+
+def test_black_pyi_file():
+    """Test consistent code formatting between isort and black for `.pyi` files.
+
+    black only allows no more than two consecutive blank lines in a `.pyi` file.
+    """
+
+    black_test(
+        """# comment
+
+
+import math
+
+from typing import Sequence
+import numpy as np
+
+
+def add(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    ...
+
+
+def sub(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    ...
+""",
+        """# comment
+
+
+import math
+from typing import Sequence
+
+import numpy as np
+
+
+def add(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    ...
+
+
+def sub(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    ...
+""",
+        is_pyi=False,
+        lines_before_imports=2,
+        lines_after_imports=2,
+    )
+
+    black_test(
+        """# comment
+
+
+import math
+
+from typing import Sequence
+import numpy as np
+
+
+def add(a: np.ndarray, b: np.ndarray) -> np.ndarray: ...
+def sub(a: np.ndarray, b: np.ndarray) -> np.ndarray: ...
+""",
+        """# comment
+
+import math
+from typing import Sequence
+
+import numpy as np
+
+def add(a: np.ndarray, b: np.ndarray) -> np.ndarray: ...
+def sub(a: np.ndarray, b: np.ndarray) -> np.ndarray: ...
+""",
+        is_pyi=True,
+        lines_before_imports=2,  # will be ignored
+        lines_after_imports=2,  # will be ignored
+    )