[ty] Add tests for relative import completions

crates/ty_ide/src/completion.rs

@@ -2129,6 +2129,57 @@ from sys import ??, <CURSOR>, ??
         test.assert_completions_include("getsizeof");
     }
 
+    #[test]
+    fn relative_from_import1() {
+        let test = CursorTest::builder()
+            .source("package/__init__.py", "")
+            .source(
+                "package/foo.py",
+                "\
+Cheetah = 1
+Lion = 2
+Cougar = 3
+",
+            )
+            .source("package/sub1/sub2/bar.py", "from ...foo import <CURSOR>")
+            .build();
+        test.assert_completions_include("Cheetah");
+    }
+
+    #[test]
+    fn relative_from_import2() {
+        let test = CursorTest::builder()
+            .source("package/__init__.py", "")
+            .source(
+                "package/sub1/foo.py",
+                "\
+Cheetah = 1
+Lion = 2
+Cougar = 3
+",
+            )
+            .source("package/sub1/sub2/bar.py", "from ..foo import <CURSOR>")
+            .build();
+        test.assert_completions_include("Cheetah");
+    }
+
+    #[test]
+    fn relative_from_import3() {
+        let test = CursorTest::builder()
+            .source("package/__init__.py", "")
+            .source(
+                "package/sub1/sub2/foo.py",
+                "\
+Cheetah = 1
+Lion = 2
+Cougar = 3
+",
+            )
+            .source("package/sub1/sub2/bar.py", "from .foo import <CURSOR>")
+            .build();
+        test.assert_completions_include("Cheetah");
+    }
+
     #[test]
     fn import_submodule_not_attribute1() {
         let test = cursor_test(
@@ -2184,7 +2235,7 @@ importlib.<CURSOR>
         }
 
         fn completions_if(&self, predicate: impl Fn(&str) -> bool) -> String {
-            let completions = completion(&self.db, self.file, self.cursor_offset);
+            let completions = completion(&self.db, self.cursor.file, self.cursor.offset);
             if completions.is_empty() {
                 return "<No completions found>".to_string();
             }
@@ -2198,7 +2249,7 @@ importlib.<CURSOR>
 
         #[track_caller]
         fn assert_completions_include(&self, expected: &str) {
-            let completions = completion(&self.db, self.file, self.cursor_offset);
+            let completions = completion(&self.db, self.cursor.file, self.cursor.offset);
 
             assert!(
                 completions
@@ -2210,7 +2261,7 @@ importlib.<CURSOR>
 
         #[track_caller]
         fn assert_completions_do_not_include(&self, unexpected: &str) {
-            let completions = completion(&self.db, self.file, self.cursor_offset);
+            let completions = completion(&self.db, self.cursor.file, self.cursor.offset);
 
             assert!(
                 completions

crates/ty_ide/src/goto.rs

@@ -833,7 +833,8 @@ f(**kwargs<CURSOR>)
 
     impl CursorTest {
         fn goto_type_definition(&self) -> String {
-            let Some(targets) = goto_type_definition(&self.db, self.file, self.cursor_offset)
+            let Some(targets) =
+                goto_type_definition(&self.db, self.cursor.file, self.cursor.offset)
             else {
                 return "No goto target found".to_string();
             };

crates/ty_ide/src/hover.rs

@@ -737,7 +737,7 @@ mod tests {
         fn hover(&self) -> String {
             use std::fmt::Write;
 
-            let Some(hover) = hover(&self.db, self.file, self.cursor_offset) else {
+            let Some(hover) = hover(&self.db, self.cursor.file, self.cursor.offset) else {
                 return "Hover provided no content".to_string();
             };
 
@@ -769,7 +769,7 @@ mod tests {
             );
             diagnostic.annotate(
                 Annotation::secondary(
-                    Span::from(source.file()).with_range(TextRange::empty(self.cursor_offset)),
+                    Span::from(source.file()).with_range(TextRange::empty(self.cursor.offset)),
                 )
                 .message("Cursor offset"),
             );

crates/ty_ide/src/lib.rs

@@ -213,59 +213,25 @@ mod tests {
         SearchPathSettings,
     };
 
+    /// A way to create a simple single-file (named `main.py`) cursor test.
+    ///
+    /// Use cases that require multiple files with a `<CURSOR>` marker
+    /// in a file other than `main.py` can use `CursorTest::builder()`.
     pub(super) fn cursor_test(source: &str) -> CursorTest {
-        let mut db = TestDb::new();
-        let cursor_offset = source.find("<CURSOR>").expect(
-            "`source`` should contain a `<CURSOR>` marker, indicating the position of the cursor.",
-        );
-
-        let mut content = source[..cursor_offset].to_string();
-        content.push_str(&source[cursor_offset + "<CURSOR>".len()..]);
-
-        db.write_file("main.py", &content)
-            .expect("write to memory file system to be successful");
-
-        let file = system_path_to_file(&db, "main.py").expect("newly written file to existing");
-
-        Program::from_settings(
-            &db,
-            ProgramSettings {
-                python_version: Some(PythonVersionWithSource::default()),
-                python_platform: PythonPlatform::default(),
-                search_paths: SearchPathSettings {
-                    extra_paths: vec![],
-                    src_roots: vec![SystemPathBuf::from("/")],
-                    custom_typeshed: None,
-                    python_path: PythonPath::KnownSitePackages(vec![]),
-                },
-            },
-        )
-        .expect("Default settings to be valid");
-
-        let mut insta_settings = insta::Settings::clone_current();
-        insta_settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
-        // Filter out TODO types because they are different between debug and release builds.
-        insta_settings.add_filter(r"@Todo\(.+\)", "@Todo");
-
-        let insta_settings_guard = insta_settings.bind_to_scope();
-
-        CursorTest {
-            db,
-            cursor_offset: TextSize::try_from(cursor_offset)
-                .expect("source to be smaller than 4GB"),
-            file,
-            _insta_settings_guard: insta_settings_guard,
-        }
+        CursorTest::builder().source("main.py", source).build()
     }
 
     pub(super) struct CursorTest {
         pub(super) db: TestDb,
-        pub(super) cursor_offset: TextSize,
-        pub(super) file: File,
+        pub(super) cursor: Cursor,
         _insta_settings_guard: SettingsBindDropGuard,
     }
 
     impl CursorTest {
+        pub(super) fn builder() -> CursorTestBuilder {
+            CursorTestBuilder::default()
+        }
+
         pub(super) fn write_file(
             &mut self,
             path: impl AsRef<SystemPath>,
@@ -295,6 +261,123 @@ mod tests {
         }
     }
 
+    /// The file and offset into that file containing
+    /// a `<CURSOR>` marker.
+    pub(super) struct Cursor {
+        pub(super) file: File,
+        pub(super) offset: TextSize,
+    }
+
+    #[derive(Default)]
+    pub(super) struct CursorTestBuilder {
+        /// A list of source files, corresponding to the
+        /// file's path and its contents.
+        sources: Vec<Source>,
+    }
+
+    impl CursorTestBuilder {
+        pub(super) fn build(&self) -> CursorTest {
+            let mut db = TestDb::new();
+            let mut cursor: Option<Cursor> = None;
+            for &Source {
+                ref path,
+                ref contents,
+                cursor_offset,
+            } in &self.sources
+            {
+                db.write_file(path, contents)
+                    .expect("write to memory file system to be successful");
+                let Some(offset) = cursor_offset else {
+                    continue;
+                };
+
+                let file = system_path_to_file(&db, path).expect("newly written file to existing");
+                // This assert should generally never trip, since
+                // we have an assert on `CursorTestBuilder::source`
+                // to ensure we never have more than one marker.
+                assert!(
+                    cursor.is_none(),
+                    "found more than one source that contains `<CURSOR>`"
+                );
+                cursor = Some(Cursor { file, offset });
+            }
+
+            Program::from_settings(
+                &db,
+                ProgramSettings {
+                    python_version: Some(PythonVersionWithSource::default()),
+                    python_platform: PythonPlatform::default(),
+                    search_paths: SearchPathSettings {
+                        extra_paths: vec![],
+                        src_roots: vec![SystemPathBuf::from("/")],
+                        custom_typeshed: None,
+                        python_path: PythonPath::KnownSitePackages(vec![]),
+                    },
+                },
+            )
+            .expect("Default settings to be valid");
+
+            let mut insta_settings = insta::Settings::clone_current();
+            insta_settings.add_filter(r#"\\(\w\w|\s|\.|")"#, "/$1");
+            // Filter out TODO types because they are different between debug and release builds.
+            insta_settings.add_filter(r"@Todo\(.+\)", "@Todo");
+
+            let insta_settings_guard = insta_settings.bind_to_scope();
+
+            CursorTest {
+                db,
+                cursor: cursor.expect("at least one source to contain `<CURSOR>`"),
+                _insta_settings_guard: insta_settings_guard,
+            }
+        }
+
+        pub(super) fn source(
+            &mut self,
+            path: impl Into<SystemPathBuf>,
+            contents: impl Into<String>,
+        ) -> &mut CursorTestBuilder {
+            const MARKER: &str = "<CURSOR>";
+
+            let path = path.into();
+            let contents = contents.into();
+            let Some(cursor_offset) = contents.find(MARKER) else {
+                self.sources.push(Source {
+                    path,
+                    contents,
+                    cursor_offset: None,
+                });
+                return self;
+            };
+
+            if let Some(source) = self.sources.iter().find(|src| src.cursor_offset.is_some()) {
+                panic!(
+                    "cursor tests must contain exactly one file \
+                     with a `<CURSOR>` marker, but found a marker \
+                     in both `{path1}` and `{path2}`",
+                    path1 = source.path,
+                    path2 = path,
+                );
+            }
+
+            let mut without_cursor_marker = contents[..cursor_offset].to_string();
+            without_cursor_marker.push_str(&contents[cursor_offset + MARKER.len()..]);
+            let cursor_offset =
+                TextSize::try_from(cursor_offset).expect("source to be smaller than 4GB");
+            self.sources.push(Source {
+                path,
+                contents: without_cursor_marker,
+                cursor_offset: Some(cursor_offset),
+            });
+            self
+        }
+    }
+
+    struct Source {
+        path: SystemPathBuf,
+        contents: String,
+        cursor_offset: Option<TextSize>,
+    }
+
     pub(super) trait IntoDiagnostic {
         fn into_diagnostic(self) -> Diagnostic;
     }