[ty] Add redeclaration LSP tests (#21812)

crates/ty_ide/src/doc_highlights.rs

@@ -236,4 +236,52 @@ def test():
 
         assert_snapshot!(test.document_highlights(), @"No highlights found");
     }
+
+    // TODO: Should only highlight the last use and the last declaration
+    #[test]
+    fn redeclarations() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                a: str = "test"
+
+                a: int = 10
+
+                print(a<CURSOR>)
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.document_highlights(), @r#"
+        info[document_highlights]: Highlight 1 (Write)
+         --> main.py:2:1
+          |
+        2 | a: str = "test"
+          | ^
+        3 |
+        4 | a: int = 10
+          |
+
+        info[document_highlights]: Highlight 2 (Write)
+         --> main.py:4:1
+          |
+        2 | a: str = "test"
+        3 |
+        4 | a: int = 10
+          | ^
+        5 |
+        6 | print(a)
+          |
+
+        info[document_highlights]: Highlight 3 (Read)
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+          |
+        "#);
+    }
 }

crates/ty_ide/src/find_references.rs

@@ -2113,4 +2113,52 @@ func<CURSOR>_alias()
           |
         ");
     }
+
+    // TODO: Should only return references to the last declaration
+    #[test]
+    fn declarations() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                a: str = "test"
+
+                a: int = 10
+
+                print(a<CURSOR>)
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.references(), @r#"
+        info[references]: Reference 1
+         --> main.py:2:1
+          |
+        2 | a: str = "test"
+          | ^
+        3 |
+        4 | a: int = 10
+          |
+
+        info[references]: Reference 2
+         --> main.py:4:1
+          |
+        2 | a: str = "test"
+        3 |
+        4 | a: int = 10
+          | ^
+        5 |
+        6 | print(a)
+          |
+
+        info[references]: Reference 3
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+          |
+        "#);
+    }
 }

crates/ty_ide/src/goto_declaration.rs

@@ -2894,6 +2894,86 @@ def ab(a: int, *, c: int): ...
         ");
     }
 
+    // TODO: Should only return `a: int`
+    #[test]
+    fn redeclarations() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                a: str = "test"
+
+                a: int = 10
+
+                print(a<CURSOR>)
+
+                a: bool = True
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.goto_declaration(), @r#"
+        info[goto-declaration]: Declaration
+         --> main.py:2:1
+          |
+        2 | a: str = "test"
+          | ^
+        3 |
+        4 | a: int = 10
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+
+        info[goto-declaration]: Declaration
+         --> main.py:4:1
+          |
+        2 | a: str = "test"
+        3 |
+        4 | a: int = 10
+          | ^
+        5 |
+        6 | print(a)
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+
+        info[goto-declaration]: Declaration
+         --> main.py:8:1
+          |
+        6 | print(a)
+        7 |
+        8 | a: bool = True
+          | ^
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+        "#);
+    }
+
     impl CursorTest {
         fn goto_declaration(&self) -> String {
             let Some(targets) = goto_declaration(&self.db, self.cursor.file, self.cursor.offset)

crates/ty_ide/src/goto_definition.rs

@@ -1714,6 +1714,86 @@ Traceb<CURSOR>ackType
         assert_snapshot!(test.goto_definition(), @"No goto target found");
     }
 
+    // TODO: Should only list `a: int`
+    #[test]
+    fn redeclarations() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                a: str = "test"
+
+                a: int = 10
+
+                print(a<CURSOR>)
+
+                a: bool = True
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.goto_definition(), @r#"
+        info[goto-definition]: Definition
+         --> main.py:2:1
+          |
+        2 | a: str = "test"
+          | ^
+        3 |
+        4 | a: int = 10
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+
+        info[goto-definition]: Definition
+         --> main.py:4:1
+          |
+        2 | a: str = "test"
+        3 |
+        4 | a: int = 10
+          | ^
+        5 |
+        6 | print(a)
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+
+        info[goto-definition]: Definition
+         --> main.py:8:1
+          |
+        6 | print(a)
+        7 |
+        8 | a: bool = True
+          | ^
+          |
+        info: Source
+         --> main.py:6:7
+          |
+        4 | a: int = 10
+        5 |
+        6 | print(a)
+          |       ^
+        7 |
+        8 | a: bool = True
+          |
+        "#);
+    }
+
     impl CursorTest {
         fn goto_definition(&self) -> String {
             let Some(targets) = goto_definition(&self.db, self.cursor.file, self.cursor.offset)

crates/ty_ide/src/rename.rs

@@ -1427,12 +1427,11 @@ result = func(10, y=20)
         ");
     }
 
-    // TODO: This should rename all overloads
     #[test]
     fn rename_overloaded_function() {
         let test = CursorTest::builder()
             .source(
-                "lib1.py",
+                "lib.py",
                 r#"
                 from typing import overload, Any
 
@@ -1450,16 +1449,16 @@ result = func(10, y=20)
             .source(
                 "main.py",
                 r#"
-                from lib2 import test
+                from lib import test
 
                 test("test")
                 "#,
             )
             .build();
 
-        assert_snapshot!(test.rename("better_name"), @r"
+        assert_snapshot!(test.rename("better_name"), @r#"
-        info[rename]: Rename symbol (found 1 locations)
+        info[rename]: Rename symbol (found 3 locations)
-         --> lib1.py:5:5
+         --> lib.py:5:5
           |
         4 | @overload
         5 | def test() -> None: ...
@@ -1467,7 +1466,117 @@ result = func(10, y=20)
         6 | @overload
         7 | def test(a: str) -> str: ...
           |
-        ");
+         ::: main.py:2:17
+          |
+        2 | from lib import test
+          |                 ----
+        3 |
+        4 | test("test")
+          | ----
+          |
+        "#);
+    }
+
+    #[test]
+    fn rename_overloaded_method() {
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+
+                class Test:
+                    @overload
+                    def test<CURSOR>() -> None: ...
+                    @overload
+                    def test(a: str) -> str: ...
+                    @overload
+                    def test(a: int) -> int: ...
+
+                    def test(a: Any) -> Any:
+                        return a
+
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import Test
+
+                Test().test("test")
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r#"
+        info[rename]: Rename symbol (found 2 locations)
+         --> lib.py:6:9
+          |
+        4 | class Test:
+        5 |     @overload
+        6 |     def test() -> None: ...
+          |         ^^^^
+        7 |     @overload
+        8 |     def test(a: str) -> str: ...
+          |
+         ::: main.py:4:8
+          |
+        2 | from lib import Test
+        3 |
+        4 | Test().test("test")
+          |        ----
+          |
+        "#);
+    }
+
+    #[test]
+    fn rename_overloaded_function_usage() {
+        let test = CursorTest::builder()
+            .source(
+                "lib.py",
+                r#"
+                from typing import overload, Any
+
+                @overload
+                def test() -> None: ...
+                @overload
+                def test(a: str) -> str: ...
+                @overload
+                def test(a: int) -> int: ...
+
+                def test(a: Any) -> Any:
+                    return a
+                "#,
+            )
+            .source(
+                "main.py",
+                r#"
+                from lib import test
+
+                test<CURSOR>("test")
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r#"
+        info[rename]: Rename symbol (found 3 locations)
+         --> main.py:2:17
+          |
+        2 | from lib import test
+          |                 ^^^^
+        3 |
+        4 | test("test")
+          | ----
+          |
+         ::: lib.py:5:5
+          |
+        4 | @overload
+        5 | def test() -> None: ...
+          |     ----
+        6 | @overload
+        7 | def test(a: str) -> str: ...
+          |
+        "#);
     }
 
     #[test]
@@ -2034,4 +2143,36 @@ result = func(10, y=20)
           |
         ");
     }
+
+    // TODO: Should not rename the first declaration
+    #[test]
+    fn rename_redeclarations() {
+        let test = CursorTest::builder()
+            .source(
+                "main.py",
+                r#"
+                a: str = "test"
+
+                a: int = 10
+
+                print(a<CURSOR>)
+                "#,
+            )
+            .build();
+
+        assert_snapshot!(test.rename("better_name"), @r#"
+        info[rename]: Rename symbol (found 3 locations)
+         --> main.py:2:1
+          |
+        2 | a: str = "test"
+          | ^
+        3 |
+        4 | a: int = 10
+          | -
+        5 |
+        6 | print(a)
+          |       -
+          |
+        "#);
+    }
 }