[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.

Author: sharkdp

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<()> {