[red-knot] Unify setup_db() functions, add TestDb builder (#14777)

## Summary

- Instead of seven (more or less similar) `setup_db` functions, use just
one in a single central place.
- For every test that needs customization beyond that, offer a
`TestDbBuilder` that can control the Python target version, custom
typeshed, and pre-existing files.

The main motivation for this is that we're soon going to need
customization of the Python version, and I didn't feel like adding this
to each of the existing `setup_db` functions.

crates/red_knot_python_semantic/src/db.rs

@@ -11,8 +11,13 @@ pub trait Db: SourceDb + Upcast<dyn SourceDb> {
 pub(crate) mod tests {
     use std::sync::Arc;
 
+    use crate::program::{Program, SearchPathSettings};
+    use crate::python_version::PythonVersion;
+    use crate::ProgramSettings;
+
+    use anyhow::Context;
     use ruff_db::files::{File, Files};
-    use ruff_db::system::{DbWithTestSystem, System, TestSystem};
+    use ruff_db::system::{DbWithTestSystem, System, SystemPathBuf, TestSystem};
     use ruff_db::vendored::VendoredFileSystem;
     use ruff_db::{Db as SourceDb, Upcast};
 
@@ -108,4 +113,66 @@ pub(crate) mod tests {
             events.push(event);
         }
     }
+
+    pub(crate) struct TestDbBuilder<'a> {
+        /// Target Python version
+        python_version: PythonVersion,
+        /// Path to a custom typeshed directory
+        custom_typeshed: Option<SystemPathBuf>,
+        /// Path and content pairs for files that should be present
+        files: Vec<(&'a str, &'a str)>,
+    }
+
+    impl<'a> TestDbBuilder<'a> {
+        pub(crate) fn new() -> Self {
+            Self {
+                python_version: PythonVersion::default(),
+                custom_typeshed: None,
+                files: vec![],
+            }
+        }
+
+        pub(crate) fn with_python_version(mut self, version: PythonVersion) -> Self {
+            self.python_version = version;
+            self
+        }
+
+        pub(crate) fn with_custom_typeshed(mut self, path: &str) -> Self {
+            self.custom_typeshed = Some(SystemPathBuf::from(path));
+            self
+        }
+
+        pub(crate) fn with_file(mut self, path: &'a str, content: &'a str) -> Self {
+            self.files.push((path, content));
+            self
+        }
+
+        pub(crate) fn build(self) -> anyhow::Result<TestDb> {
+            let mut db = TestDb::new();
+
+            let src_root = SystemPathBuf::from("/src");
+            db.memory_file_system().create_directory_all(&src_root)?;
+
+            db.write_files(self.files)
+                .context("Failed to write test files")?;
+
+            let mut search_paths = SearchPathSettings::new(src_root);
+            search_paths.custom_typeshed = self.custom_typeshed;
+
+            Program::from_settings(
+                &db,
+                &ProgramSettings {
+                    target_version: self.python_version,
+                    search_paths,
+                },
+            )
+            .context("Failed to configure Program settings")?;
+
+            Ok(db)
+        }
+    }
+
+    pub(crate) fn setup_db() -> TestDb {
+        TestDbBuilder::new().build().expect("valid TestDb setup")
+    }
 }

crates/red_knot_python_semantic/src/semantic_model.rs

@@ -166,31 +166,15 @@ impl_binding_has_ty!(ast::ParameterWithDefault);
 mod tests {
     use ruff_db::files::system_path_to_file;
     use ruff_db::parsed::parsed_module;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
 
-    use crate::db::tests::TestDb;
-    use crate::program::{Program, SearchPathSettings};
-    use crate::python_version::PythonVersion;
-    use crate::{HasTy, ProgramSettings, SemanticModel};
-
-    fn setup_db<'a>(files: impl IntoIterator<Item = (&'a str, &'a str)>) -> anyhow::Result<TestDb> {
-        let mut db = TestDb::new();
-        db.write_files(files)?;
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(SystemPathBuf::from("/src")),
-            },
-        )?;
-
-        Ok(db)
-    }
+    use crate::db::tests::TestDbBuilder;
+    use crate::{HasTy, SemanticModel};
 
     #[test]
     fn function_ty() -> anyhow::Result<()> {
-        let db = setup_db([("/src/foo.py", "def test(): pass")])?;
+        let db = TestDbBuilder::new()
+            .with_file("/src/foo.py", "def test(): pass")
+            .build()?;
 
         let foo = system_path_to_file(&db, "/src/foo.py").unwrap();
 
@@ -207,7 +191,9 @@ mod tests {
 
     #[test]
     fn class_ty() -> anyhow::Result<()> {
-        let db = setup_db([("/src/foo.py", "class Test: pass")])?;
+        let db = TestDbBuilder::new()
+            .with_file("/src/foo.py", "class Test: pass")
+            .build()?;
 
         let foo = system_path_to_file(&db, "/src/foo.py").unwrap();
 
@@ -224,10 +210,10 @@ mod tests {
 
     #[test]
     fn alias_ty() -> anyhow::Result<()> {
-        let db = setup_db([
-            ("/src/foo.py", "class Test: pass"),
-            ("/src/bar.py", "from foo import Test"),
-        ])?;
+        let db = TestDbBuilder::new()
+            .with_file("/src/foo.py", "class Test: pass")
+            .with_file("/src/bar.py", "from foo import Test")
+            .build()?;
 
         let bar = system_path_to_file(&db, "/src/bar.py").unwrap();
 

crates/red_knot_python_semantic/src/symbol.rs

@@ -90,7 +90,7 @@ impl<'db> Symbol<'db> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::types::tests::setup_db;
+    use crate::db::tests::setup_db;
 
     #[test]
     fn test_symbol_or_fall_back_to() {

crates/red_knot_python_semantic/src/types.rs

@@ -3203,38 +3203,16 @@ static_assertions::assert_eq_size!(Type, [u8; 16]);
 #[cfg(test)]
 pub(crate) mod tests {
     use super::*;
-    use crate::db::tests::TestDb;
-    use crate::program::{Program, SearchPathSettings};
-    use crate::python_version::PythonVersion;
+    use crate::db::tests::{setup_db, TestDb, TestDbBuilder};
     use crate::stdlib::typing_symbol;
-    use crate::ProgramSettings;
+    use crate::PythonVersion;
     use ruff_db::files::system_path_to_file;
     use ruff_db::parsed::parsed_module;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+    use ruff_db::system::DbWithTestSystem;
     use ruff_db::testing::assert_function_query_was_not_run;
     use ruff_python_ast as ast;
     use test_case::test_case;
 
-    pub(crate) fn setup_db() -> TestDb {
-        let db = TestDb::new();
-
-        let src_root = SystemPathBuf::from("/src");
-        db.memory_file_system()
-            .create_directory_all(&src_root)
-            .unwrap();
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(src_root),
-            },
-        )
-        .expect("Valid search path settings");
-
-        db
-    }
-
     /// A test representation of a type that can be transformed unambiguously into a real Type,
     /// given a db.
     #[derive(Debug, Clone, PartialEq)]
@@ -3839,7 +3817,10 @@ pub(crate) mod tests {
 
     #[test]
     fn typing_vs_typeshed_no_default() {
-        let db = setup_db();
+        let db = TestDbBuilder::new()
+            .with_python_version(PythonVersion::PY313)
+            .build()
+            .unwrap();
 
         let typing_no_default = typing_symbol(&db, "NoDefault").expect_type();
         let typing_extensions_no_default = typing_extensions_symbol(&db, "NoDefault").expect_type();

crates/red_knot_python_semantic/src/types/builder.rs

@@ -378,35 +378,14 @@ impl<'db> InnerIntersectionBuilder<'db> {
 #[cfg(test)]
 mod tests {
     use super::{IntersectionBuilder, IntersectionType, Type, UnionType};
-    use crate::db::tests::TestDb;
-    use crate::program::{Program, SearchPathSettings};
-    use crate::python_version::PythonVersion;
+
+    use crate::db::tests::{setup_db, TestDb};
     use crate::types::{global_symbol, todo_type, KnownClass, UnionBuilder};
-    use crate::ProgramSettings;
+
     use ruff_db::files::system_path_to_file;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+    use ruff_db::system::DbWithTestSystem;
     use test_case::test_case;
 
-    fn setup_db() -> TestDb {
-        let db = TestDb::new();
-
-        let src_root = SystemPathBuf::from("/src");
-        db.memory_file_system()
-            .create_directory_all(&src_root)
-            .unwrap();
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(src_root),
-            },
-        )
-        .expect("Valid search path settings");
-
-        db
-    }
-
     #[test]
     fn build_union() {
         let db = setup_db();

crates/red_knot_python_semantic/src/types/display.rs

@@ -357,31 +357,10 @@ impl Display for DisplayStringLiteralType<'_> {
 #[cfg(test)]
 mod tests {
     use ruff_db::files::system_path_to_file;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+    use ruff_db::system::DbWithTestSystem;
 
-    use crate::db::tests::TestDb;
+    use crate::db::tests::setup_db;
     use crate::types::{global_symbol, SliceLiteralType, StringLiteralType, Type, UnionType};
-    use crate::{Program, ProgramSettings, PythonVersion, SearchPathSettings};
-
-    fn setup_db() -> TestDb {
-        let db = TestDb::new();
-
-        let src_root = SystemPathBuf::from("/src");
-        db.memory_file_system()
-            .create_directory_all(&src_root)
-            .unwrap();
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(src_root),
-            },
-        )
-        .expect("Valid search path settings");
-
-        db
-    }
 
     #[test]
     fn test_condense_literal_display_by_type() -> anyhow::Result<()> {

crates/red_knot_python_semantic/src/types/infer.rs

@@ -5046,68 +5046,19 @@ fn perform_membership_test_comparison<'db>(
 
 #[cfg(test)]
 mod tests {
-    use anyhow::Context;
-
-    use crate::db::tests::TestDb;
-    use crate::program::{Program, SearchPathSettings};
-    use crate::python_version::PythonVersion;
+    use crate::db::tests::{setup_db, TestDb, TestDbBuilder};
     use crate::semantic_index::definition::Definition;
     use crate::semantic_index::symbol::FileScopeId;
     use crate::semantic_index::{global_scope, semantic_index, symbol_table, use_def_map};
     use crate::types::check_types;
-    use crate::{HasTy, ProgramSettings, SemanticModel};
+    use crate::{HasTy, SemanticModel};
     use ruff_db::files::{system_path_to_file, File};
     use ruff_db::parsed::parsed_module;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+    use ruff_db::system::DbWithTestSystem;
     use ruff_db::testing::assert_function_query_was_not_run;
 
     use super::*;
 
-    fn setup_db() -> TestDb {
-        let db = TestDb::new();
-
-        let src_root = SystemPathBuf::from("/src");
-        db.memory_file_system()
-            .create_directory_all(&src_root)
-            .unwrap();
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(src_root),
-            },
-        )
-        .expect("Valid search path settings");
-
-        db
-    }
-
-    fn setup_db_with_custom_typeshed<'a>(
-        typeshed: &str,
-        files: impl IntoIterator<Item = (&'a str, &'a str)>,
-    ) -> anyhow::Result<TestDb> {
-        let mut db = TestDb::new();
-        let src_root = SystemPathBuf::from("/src");
-
-        db.write_files(files)
-            .context("Failed to write test files")?;
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings {
-                    custom_typeshed: Some(SystemPathBuf::from(typeshed)),
-                    ..SearchPathSettings::new(src_root)
-                },
-            },
-        )
-        .context("Failed to create Program")?;
-
-        Ok(db)
-    }
-
     #[track_caller]
     fn assert_public_ty(db: &TestDb, file_name: &str, symbol_name: &str, expected: &str) {
         let file = system_path_to_file(db, file_name).expect("file to exist");
@@ -5502,17 +5453,15 @@ mod tests {
 
     #[test]
     fn builtin_symbol_custom_stdlib() -> anyhow::Result<()> {
-        let db = setup_db_with_custom_typeshed(
-            "/typeshed",
-            [
-                ("/src/a.py", "c = copyright"),
-                (
-                    "/typeshed/stdlib/builtins.pyi",
-                    "def copyright() -> None: ...",
-                ),
-                ("/typeshed/stdlib/VERSIONS", "builtins: 3.8-"),
-            ],
-        )?;
+        let db = TestDbBuilder::new()
+            .with_custom_typeshed("/typeshed")
+            .with_file("/src/a.py", "c = copyright")
+            .with_file(
+                "/typeshed/stdlib/builtins.pyi",
+                "def copyright() -> None: ...",
+            )
+            .with_file("/typeshed/stdlib/VERSIONS", "builtins: 3.8-")
+            .build()?;
 
         assert_public_ty(&db, "/src/a.py", "c", "Literal[copyright]");
 
@@ -5521,14 +5470,12 @@ mod tests {
 
     #[test]
     fn unknown_builtin_later_defined() -> anyhow::Result<()> {
-        let db = setup_db_with_custom_typeshed(
-            "/typeshed",
-            [
-                ("/src/a.py", "x = foo"),
-                ("/typeshed/stdlib/builtins.pyi", "foo = bar; bar = 1"),
-                ("/typeshed/stdlib/VERSIONS", "builtins: 3.8-"),
-            ],
-        )?;
+        let db = TestDbBuilder::new()
+            .with_custom_typeshed("/typeshed")
+            .with_file("/src/a.py", "x = foo")
+            .with_file("/typeshed/stdlib/builtins.pyi", "foo = bar; bar = 1")
+            .with_file("/typeshed/stdlib/VERSIONS", "builtins: 3.8-")
+            .build()?;
 
         assert_public_ty(&db, "/src/a.py", "x", "Unknown");
 

crates/red_knot_python_semantic/src/types/property_tests.rs

@@ -26,8 +26,8 @@
 
 use std::sync::{Arc, Mutex, MutexGuard, OnceLock};
 
-use super::tests::{setup_db, Ty};
-use crate::db::tests::TestDb;
+use super::tests::Ty;
+use crate::db::tests::{setup_db, TestDb};
 use crate::types::KnownClass;
 use quickcheck::{Arbitrary, Gen};
 

crates/red_knot_python_semantic/src/types/signatures.rs

@@ -189,32 +189,9 @@ impl<'db> Parameter<'db> {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::db::tests::TestDb;
-    use crate::program::{Program, SearchPathSettings};
-    use crate::python_version::PythonVersion;
+    use crate::db::tests::{setup_db, TestDb};
     use crate::types::{global_symbol, FunctionType};
-    use crate::ProgramSettings;
-    use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
-
-    pub(crate) fn setup_db() -> TestDb {
-        let db = TestDb::new();
-
-        let src_root = SystemPathBuf::from("/src");
-        db.memory_file_system()
-            .create_directory_all(&src_root)
-            .unwrap();
-
-        Program::from_settings(
-            &db,
-            &ProgramSettings {
-                target_version: PythonVersion::default(),
-                search_paths: SearchPathSettings::new(src_root),
-            },
-        )
-        .expect("Valid search path settings");
-
-        db
-    }
+    use ruff_db::system::DbWithTestSystem;
 
     #[track_caller]
     fn get_function_f<'db>(db: &'db TestDb, file: &'static str) -> FunctionType<'db> {