[ty] Add a cursor test builder

This doesn't change any functionality of the cursor tests, but does
re-arrange the code a bit. Firstly, it's now in a builder. And secondly,
there's an API to add multiple files to the test (but exactly one must
have a `<CURSOR>` marker).

crates/ty_ide/src/lib.rs

@@ -214,48 +214,7 @@ mod tests {
     };
 
     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 {
@@ -266,6 +225,10 @@ mod tests {
     }
 
     impl CursorTest {
+        pub(super) fn builder() -> CursorTestBuilder {
+            CursorTestBuilder::default()
+        }
+
         pub(super) fn write_file(
             &mut self,
             path: impl AsRef<SystemPath>,
@@ -295,6 +258,124 @@ mod tests {
         }
     }
 
+    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();
+            let Cursor { file, offset } =
+                cursor.expect("at least one source to contain `<CURSOR>`");
+
+            CursorTest {
+                db,
+                cursor_offset: offset,
+                file,
+                _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;
     }