[red-knot] Add more stress tests for module resolver invalidation (#1…

…2272)

crates/red_knot_module_resolver/src/resolver.rs

@@ -384,8 +384,10 @@ impl PackageKind {
 
 #[cfg(test)]
 mod tests {
+    use internal::ModuleNameIngredient;
     use ruff_db::files::{system_path_to_file, File, FilePath};
     use ruff_db::system::{DbWithTestSystem, OsSystem, SystemPath};
+    use ruff_db::testing::assert_function_query_was_not_run;
 
     use crate::db::tests::TestDb;
     use crate::module::ModuleKind;
@@ -966,7 +968,7 @@ mod tests {
     }
 
     #[test]
-    fn adding_a_file_on_which_the_module_resolution_depends_on_invalidates_the_query(
+    fn adding_file_on_which_module_resolution_depends_invalidates_previously_failing_query_that_now_succeeds(
     ) -> anyhow::Result<()> {
         let TestCase { mut db, src, .. } = TestCaseBuilder::new().build();
         let foo_path = src.join("foo.py");
@@ -986,7 +988,7 @@ mod tests {
     }
 
     #[test]
-    fn removing_a_file_that_the_module_resolution_depends_on_invalidates_the_query(
+    fn removing_file_on_which_module_resolution_depends_invalidates_previously_successful_query_that_now_fails(
     ) -> anyhow::Result<()> {
         const SRC: &[FileSpec] = &[("foo.py", "x = 1"), ("foo/__init__.py", "x = 2")];
 
@@ -1009,4 +1011,133 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    fn adding_file_to_search_path_with_lower_priority_does_not_invalidate_query() {
+        const TYPESHED: MockedTypeshed = MockedTypeshed {
+            versions: "functools: 3.8-",
+            stdlib_files: &[("functools.pyi", "def update_wrapper(): ...")],
+        };
+
+        let TestCase {
+            mut db,
+            stdlib,
+            site_packages,
+            ..
+        } = TestCaseBuilder::new()
+            .with_custom_typeshed(TYPESHED)
+            .with_target_version(TargetVersion::Py38)
+            .build();
+
+        let functools_module_name = ModuleName::new_static("functools").unwrap();
+        let stdlib_functools_path = stdlib.join("functools.pyi");
+
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        assert_eq!(functools_module.search_path(), stdlib);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, &stdlib_functools_path)
+        );
+
+        // Adding a file to site-packages does not invalidate the query,
+        // since site-packages takes lower priority in the module resolution
+        db.clear_salsa_events();
+        let site_packages_functools_path = site_packages.join("functools.py");
+        db.write_file(&site_packages_functools_path, "f: int")
+            .unwrap();
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        let events = db.take_salsa_events();
+        assert_function_query_was_not_run::<resolve_module_query, _, _>(
+            &db,
+            |res| &res.function,
+            &ModuleNameIngredient::new(&db, functools_module_name.clone()),
+            &events,
+        );
+        assert_eq!(functools_module.search_path(), stdlib);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, &stdlib_functools_path)
+        );
+    }
+
+    #[test]
+    fn adding_file_to_search_path_with_higher_priority_invalidates_the_query() {
+        const TYPESHED: MockedTypeshed = MockedTypeshed {
+            versions: "functools: 3.8-",
+            stdlib_files: &[("functools.pyi", "def update_wrapper(): ...")],
+        };
+
+        let TestCase {
+            mut db,
+            stdlib,
+            src,
+            ..
+        } = TestCaseBuilder::new()
+            .with_custom_typeshed(TYPESHED)
+            .with_target_version(TargetVersion::Py38)
+            .build();
+
+        let functools_module_name = ModuleName::new_static("functools").unwrap();
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        assert_eq!(functools_module.search_path(), stdlib);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, stdlib.join("functools.pyi"))
+        );
+
+        // Adding a first-party file invalidates the query,
+        // since first-party files take higher priority in module resolution:
+        let src_functools_path = src.join("functools.py");
+        db.write_file(&src_functools_path, "FOO: int").unwrap();
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        assert_eq!(functools_module.search_path(), src);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, &src_functools_path)
+        );
+    }
+
+    #[test]
+    fn deleting_file_from_higher_priority_search_path_invalidates_the_query() {
+        const SRC: &[FileSpec] = &[("functools.py", "FOO: int")];
+
+        const TYPESHED: MockedTypeshed = MockedTypeshed {
+            versions: "functools: 3.8-",
+            stdlib_files: &[("functools.pyi", "def update_wrapper(): ...")],
+        };
+
+        let TestCase {
+            mut db,
+            stdlib,
+            src,
+            ..
+        } = TestCaseBuilder::new()
+            .with_src_files(SRC)
+            .with_custom_typeshed(TYPESHED)
+            .with_target_version(TargetVersion::Py38)
+            .build();
+
+        let functools_module_name = ModuleName::new_static("functools").unwrap();
+        let src_functools_path = src.join("functools.py");
+
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        assert_eq!(functools_module.search_path(), src);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, &src_functools_path)
+        );
+
+        // If we now delete the first-party file,
+        // it should resolve to the stdlib:
+        db.memory_file_system()
+            .remove_file(&src_functools_path)
+            .unwrap();
+        File::touch_path(&mut db, &src_functools_path);
+        let functools_module = resolve_module(&db, functools_module_name.clone()).unwrap();
+        assert_eq!(functools_module.search_path(), stdlib);
+        assert_eq!(
+            Some(functools_module.file()),
+            system_path_to_file(&db, stdlib.join("functools.pyi"))
+        );
+    }
 }

crates/red_knot_python_semantic/src/db.rs

@@ -35,13 +35,8 @@ pub trait Db:
 
 #[cfg(test)]
 pub(crate) mod tests {
-    use std::fmt::Formatter;
-    use std::marker::PhantomData;
     use std::sync::Arc;
 
-    use salsa::id::AsId;
-    use salsa::ingredient::Ingredient;
-    use salsa::storage::HasIngredientsFor;
     use salsa::DebugWithDb;
 
     use red_knot_module_resolver::{vendored_typeshed_stubs, Db as ResolverDb, Jar as ResolverJar};
@@ -150,106 +145,4 @@ pub(crate) mod tests {
             })
         }
     }
-
-    pub(crate) fn assert_will_run_function_query<'db, C, Db, Jar>(
-        db: &'db Db,
-        to_function: impl FnOnce(&C) -> &salsa::function::FunctionIngredient<C>,
-        input: &C::Input<'db>,
-        events: &[salsa::Event],
-    ) where
-        C: salsa::function::Configuration<Jar = Jar>
-            + salsa::storage::IngredientsFor<Jar = Jar, Ingredients = C>,
-        Jar: HasIngredientsFor<C>,
-        Db: salsa::DbWithJar<Jar>,
-        C::Input<'db>: AsId,
-    {
-        will_run_function_query(db, to_function, input, events, true);
-    }
-
-    pub(crate) fn assert_will_not_run_function_query<'db, C, Db, Jar>(
-        db: &'db Db,
-        to_function: impl FnOnce(&C) -> &salsa::function::FunctionIngredient<C>,
-        input: &C::Input<'db>,
-        events: &[salsa::Event],
-    ) where
-        C: salsa::function::Configuration<Jar = Jar>
-            + salsa::storage::IngredientsFor<Jar = Jar, Ingredients = C>,
-        Jar: HasIngredientsFor<C>,
-        Db: salsa::DbWithJar<Jar>,
-        C::Input<'db>: AsId,
-    {
-        will_run_function_query(db, to_function, input, events, false);
-    }
-
-    fn will_run_function_query<'db, C, Db, Jar>(
-        db: &'db Db,
-        to_function: impl FnOnce(&C) -> &salsa::function::FunctionIngredient<C>,
-        input: &C::Input<'db>,
-        events: &[salsa::Event],
-        should_run: bool,
-    ) where
-        C: salsa::function::Configuration<Jar = Jar>
-            + salsa::storage::IngredientsFor<Jar = Jar, Ingredients = C>,
-        Jar: HasIngredientsFor<C>,
-        Db: salsa::DbWithJar<Jar>,
-        C::Input<'db>: AsId,
-    {
-        let (jar, _) =
-            <_ as salsa::storage::HasJar<<C as salsa::storage::IngredientsFor>::Jar>>::jar(db);
-        let ingredient = jar.ingredient();
-
-        let function_ingredient = to_function(ingredient);
-
-        let ingredient_index =
-            <salsa::function::FunctionIngredient<C> as Ingredient<Db>>::ingredient_index(
-                function_ingredient,
-            );
-
-        let did_run = events.iter().any(|event| {
-            if let salsa::EventKind::WillExecute { database_key } = event.kind {
-                database_key.ingredient_index() == ingredient_index
-                    && database_key.key_index() == input.as_id()
-            } else {
-                false
-            }
-        });
-
-        if should_run && !did_run {
-            panic!(
-                "Expected query {:?} to run but it didn't",
-                DebugIdx {
-                    db: PhantomData::<Db>,
-                    value_id: input.as_id(),
-                    ingredient: function_ingredient,
-                }
-            );
-        } else if !should_run && did_run {
-            panic!(
-                "Expected query {:?} not to run but it did",
-                DebugIdx {
-                    db: PhantomData::<Db>,
-                    value_id: input.as_id(),
-                    ingredient: function_ingredient,
-                }
-            );
-        }
-    }
-
-    struct DebugIdx<'a, I, Db>
-    where
-        I: Ingredient<Db>,
-    {
-        value_id: salsa::Id,
-        ingredient: &'a I,
-        db: PhantomData<Db>,
-    }
-
-    impl<'a, I, Db> std::fmt::Debug for DebugIdx<'a, I, Db>
-    where
-        I: Ingredient<Db>,
-    {
-        fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-            self.ingredient.fmt_index(Some(self.value_id), f)
-        }
-    }
 }

crates/red_knot_python_semantic/src/types.rs

@@ -277,10 +277,9 @@ mod tests {
     use ruff_db::files::system_path_to_file;
     use ruff_db::parsed::parsed_module;
     use ruff_db::system::{DbWithTestSystem, SystemPathBuf};
+    use ruff_db::testing::{assert_function_query_was_not_run, assert_function_query_was_run};
 
-    use crate::db::tests::{
-        assert_will_not_run_function_query, assert_will_run_function_query, TestDb,
-    };
+    use crate::db::tests::TestDb;
     use crate::semantic_index::root_scope;
     use crate::types::{infer_types, public_symbol_ty_by_name};
     use crate::{HasTy, SemanticModel};
@@ -347,7 +346,7 @@ mod tests {
         let events = db.take_salsa_events();
 
         let a_root_scope = root_scope(&db, a);
-        assert_will_run_function_query::<infer_types, _, _>(
+        assert_function_query_was_run::<infer_types, _, _>(
             &db,
             |ty| &ty.function,
             &a_root_scope,
@@ -385,7 +384,7 @@ mod tests {
 
         let a_root_scope = root_scope(&db, a);
 
-        assert_will_not_run_function_query::<infer_types, _, _>(
+        assert_function_query_was_not_run::<infer_types, _, _>(
             &db,
             |ty| &ty.function,
             &a_root_scope,
@@ -422,7 +421,7 @@ mod tests {
         let events = db.take_salsa_events();
 
         let a_root_scope = root_scope(&db, a);
-        assert_will_not_run_function_query::<infer_types, _, _>(
+        assert_function_query_was_not_run::<infer_types, _, _>(
             &db,
             |ty| &ty.function,
             &a_root_scope,

crates/ruff_db/src/lib.rs

@@ -14,6 +14,7 @@ pub mod files;
 pub mod parsed;
 pub mod source;
 pub mod system;
+pub mod testing;
 pub mod vendored;
 
 pub(crate) type FxDashMap<K, V> = dashmap::DashMap<K, V, BuildHasherDefault<FxHasher>>;