From c2e17d0399dcf6bd90f2cc82728fc8e719fe3477 Mon Sep 17 00:00:00 2001
From: Micha Reiser <micha@reiser.io>
Date: Tue, 3 Dec 2024 08:22:42 +0100
Subject: [PATCH] Possible fix for flaky file watching test (#14543)

---
 crates/red_knot/tests/file_watching.rs | 159 ++++++++++++++++---------
 crates/red_knot_workspace/src/watch.rs |  16 +++
 2 files changed, 119 insertions(+), 56 deletions(-)

diff --git a/crates/red_knot/tests/file_watching.rs b/crates/red_knot/tests/file_watching.rs
index f063bcf3d5..aeda6c5492 100644
--- a/crates/red_knot/tests/file_watching.rs
+++ b/crates/red_knot/tests/file_watching.rs
@@ -1,25 +1,23 @@
 #![allow(clippy::disallowed_names)]
 
 use std::io::Write;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use anyhow::{anyhow, Context};
 use red_knot_python_semantic::{resolve_module, ModuleName, Program, PythonVersion, SitePackages};
 use red_knot_workspace::db::{Db, RootDatabase};
-use red_knot_workspace::watch;
-use red_knot_workspace::watch::{directory_watcher, WorkspaceWatcher};
+use red_knot_workspace::watch::{directory_watcher, ChangeEvent, WorkspaceWatcher};
 use red_knot_workspace::workspace::settings::{Configuration, SearchPathConfiguration};
 use red_knot_workspace::workspace::WorkspaceMetadata;
 use ruff_db::files::{system_path_to_file, File, FileError};
 use ruff_db::source::source_text;
 use ruff_db::system::{OsSystem, SystemPath, SystemPathBuf};
-use ruff_db::testing::{setup_logging, setup_logging_with_filter};
 use ruff_db::Upcast;
 
 struct TestCase {
     db: RootDatabase,
     watcher: Option<WorkspaceWatcher>,
-    changes_receiver: crossbeam::channel::Receiver<Vec<watch::ChangeEvent>>,
+    changes_receiver: crossbeam::channel::Receiver<Vec<ChangeEvent>>,
     /// The temporary directory that contains the test files.
     /// We need to hold on to it in the test case or the temp files get deleted.
     _temp_dir: tempfile::TempDir,
@@ -40,12 +38,36 @@ impl TestCase {
         &self.db
     }
 
-    fn stop_watch(&mut self) -> Vec<watch::ChangeEvent> {
-        self.try_stop_watch(Duration::from_secs(10))
-            .expect("Expected watch changes but observed none")
+    #[track_caller]
+    fn stop_watch<M>(&mut self, matcher: M) -> Vec<ChangeEvent>
+    where
+        M: MatchEvent,
+    {
+        // track_caller is unstable for lambdas -> That's why this is a fn
+        #[track_caller]
+        fn panic_with_formatted_events(events: Vec<ChangeEvent>) -> Vec<ChangeEvent> {
+            panic!(
+                "Didn't observe expected change:\n{}",
+                events
+                    .into_iter()
+                    .map(|event| format!("  - {event:?}"))
+                    .collect::<Vec<_>>()
+                    .join("\n")
+            )
+        }
+
+        self.try_stop_watch(matcher, Duration::from_secs(10))
+            .unwrap_or_else(panic_with_formatted_events)
     }
 
-    fn try_stop_watch(&mut self, timeout: Duration) -> Option<Vec<watch::ChangeEvent>> {
+    fn try_stop_watch<M>(
+        &mut self,
+        mut matcher: M,
+        timeout: Duration,
+    ) -> Result<Vec<ChangeEvent>, Vec<ChangeEvent>>
+    where
+        M: MatchEvent,
+    {
         tracing::debug!("Try stopping watch with timeout {:?}", timeout);
 
         let watcher = self
@@ -53,36 +75,50 @@ impl TestCase {
             .take()
             .expect("Cannot call `stop_watch` more than once");
 
-        let mut all_events = self
-            .changes_receiver
-            .recv_timeout(timeout)
-            .unwrap_or_default();
+        let start = Instant::now();
+        let mut all_events = Vec::new();
+
+        loop {
+            let events = self
+                .changes_receiver
+                .recv_timeout(Duration::from_millis(100))
+                .unwrap_or_default();
+
+            if events
+                .iter()
+                .any(|event| matcher.match_event(event) || event.is_rescan())
+            {
+                all_events.extend(events);
+                break;
+            }
+
+            all_events.extend(events);
+
+            if start.elapsed() > timeout {
+                return Err(all_events);
+            }
+        }
 
         watcher.flush();
         tracing::debug!("Flushed file watcher");
         watcher.stop();
         tracing::debug!("Stopping file watcher");
 
+        // Consume remaining events
         for event in &self.changes_receiver {
             all_events.extend(event);
         }
 
-        if all_events.is_empty() {
-            return None;
-        }
-
-        Some(all_events)
+        Ok(all_events)
     }
 
-    fn take_watch_changes(&self) -> Vec<watch::ChangeEvent> {
+    fn take_watch_changes(&self) -> Vec<ChangeEvent> {
         self.try_take_watch_changes(Duration::from_secs(10))
             .expect("Expected watch changes but observed none")
     }
 
-    fn try_take_watch_changes(&self, timeout: Duration) -> Option<Vec<watch::ChangeEvent>> {
-        let Some(watcher) = &self.watcher else {
-            return None;
-        };
+    fn try_take_watch_changes(&self, timeout: Duration) -> Option<Vec<ChangeEvent>> {
+        let watcher = self.watcher.as_ref()?;
 
         let mut all_events = self
             .changes_receiver
@@ -104,7 +140,7 @@ impl TestCase {
         Some(all_events)
     }
 
-    fn apply_changes(&mut self, changes: Vec<watch::ChangeEvent>) {
+    fn apply_changes(&mut self, changes: Vec<ChangeEvent>) {
         self.db.apply_changes(changes, Some(&self.configuration));
     }
 
@@ -140,6 +176,23 @@ impl TestCase {
     }
 }
 
+trait MatchEvent {
+    fn match_event(&mut self, event: &ChangeEvent) -> bool;
+}
+
+fn event_for_file(name: &str) -> impl MatchEvent + '_ {
+    |event: &ChangeEvent| event.file_name() == Some(name)
+}
+
+impl<F> MatchEvent for F
+where
+    F: FnMut(&ChangeEvent) -> bool,
+{
+    fn match_event(&mut self, event: &ChangeEvent) -> bool {
+        (*self)(event)
+    }
+}
+
 trait SetupFiles {
     fn setup(self, root_path: &SystemPath, workspace_path: &SystemPath) -> anyhow::Result<()>;
 }
@@ -315,7 +368,7 @@ fn new_file() -> anyhow::Result<()> {
 
     std::fs::write(foo_path.as_std_path(), "print('Hello')")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -338,7 +391,7 @@ fn new_ignored_file() -> anyhow::Result<()> {
 
     std::fs::write(foo_path.as_std_path(), "print('Hello')")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -360,7 +413,7 @@ fn changed_file() -> anyhow::Result<()> {
 
     update_file(&foo_path, "print('Version 2')")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     assert!(!changes.is_empty());
 
@@ -385,7 +438,7 @@ fn deleted_file() -> anyhow::Result<()> {
 
     std::fs::remove_file(foo_path.as_std_path())?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -417,7 +470,7 @@ fn move_file_to_trash() -> anyhow::Result<()> {
         trash_path.join("foo.py").as_std_path(),
     )?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -449,7 +502,7 @@ fn move_file_to_workspace() -> anyhow::Result<()> {
 
     std::fs::rename(foo_path.as_std_path(), foo_in_workspace_path.as_std_path())?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -477,7 +530,7 @@ fn rename_file() -> anyhow::Result<()> {
 
     std::fs::rename(foo_path.as_std_path(), bar_path.as_std_path())?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("bar.py"));
 
     case.apply_changes(changes);
 
@@ -521,7 +574,7 @@ fn directory_moved_to_workspace() -> anyhow::Result<()> {
     std::fs::rename(sub_original_path.as_std_path(), sub_new_path.as_std_path())
         .with_context(|| "Failed to move sub directory")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("sub"));
 
     case.apply_changes(changes);
 
@@ -580,7 +633,7 @@ fn directory_moved_to_trash() -> anyhow::Result<()> {
     std::fs::rename(sub_path.as_std_path(), trashed_sub.as_std_path())
         .with_context(|| "Failed to move the sub directory to the trash")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("sub"));
 
     case.apply_changes(changes);
 
@@ -604,8 +657,6 @@ fn directory_moved_to_trash() -> anyhow::Result<()> {
 
 #[test]
 fn directory_renamed() -> anyhow::Result<()> {
-    let _tracing = setup_logging_with_filter("file_watching=TRACE,red_knot=TRACE");
-
     let mut case = setup([
         ("bar.py", "import sub.a"),
         ("sub/__init__.py", ""),
@@ -644,11 +695,8 @@ fn directory_renamed() -> anyhow::Result<()> {
     std::fs::rename(sub_path.as_std_path(), foo_baz.as_std_path())
         .with_context(|| "Failed to move the sub directory")?;
 
-    let changes = case.stop_watch();
-
-    for event in &changes {
-        tracing::debug!("Event: {:?}", event);
-    }
+    // Linux and windows only emit an event for the newly created root directory, but not for every new component.
+    let changes = case.stop_watch(event_for_file("sub"));
 
     case.apply_changes(changes);
 
@@ -721,7 +769,7 @@ fn directory_deleted() -> anyhow::Result<()> {
     std::fs::remove_dir_all(sub_path.as_std_path())
         .with_context(|| "Failed to remove the sub directory")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("sub"));
 
     case.apply_changes(changes);
 
@@ -758,7 +806,7 @@ fn search_path() -> anyhow::Result<()> {
 
     std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("a.py"));
 
     case.apply_changes(changes);
 
@@ -789,7 +837,7 @@ fn add_search_path() -> anyhow::Result<()> {
 
     std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("a.py"));
 
     case.apply_changes(changes);
 
@@ -818,9 +866,9 @@ fn remove_search_path() -> anyhow::Result<()> {
 
     std::fs::write(site_packages.join("a.py").as_std_path(), "class A: ...")?;
 
-    let changes = case.try_stop_watch(Duration::from_millis(100));
+    let changes = case.try_stop_watch(|_: &ChangeEvent| true, Duration::from_millis(100));
 
-    assert_eq!(changes, None);
+    assert_eq!(changes, Err(vec![]));
 
     Ok(())
 }
@@ -858,7 +906,7 @@ fn changed_versions_file() -> anyhow::Result<()> {
         "os: 3.0-",
     )?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("VERSIONS"));
 
     case.apply_changes(changes);
 
@@ -911,7 +959,7 @@ fn hard_links_in_workspace() -> anyhow::Result<()> {
     // Write to the hard link target.
     update_file(foo_path, "print('Version 2')").context("Failed to update foo.py")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("foo.py"));
 
     case.apply_changes(changes);
 
@@ -982,7 +1030,7 @@ fn hard_links_to_target_outside_workspace() -> anyhow::Result<()> {
     // Write to the hard link target.
     update_file(foo_path, "print('Version 2')").context("Failed to update foo.py")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(ChangeEvent::is_changed);
 
     case.apply_changes(changes);
 
@@ -1021,7 +1069,7 @@ mod unix {
         )
         .with_context(|| "Failed to set file permissions.")?;
 
-        let changes = case.stop_watch();
+        let changes = case.stop_watch(event_for_file("foo.py"));
 
         case.apply_changes(changes);
 
@@ -1119,7 +1167,7 @@ mod unix {
         update_file(baz_workspace, "def baz(): print('Version 3')")
             .context("Failed to update bar/baz.py")?;
 
-        let changes = case.stop_watch();
+        let changes = case.stop_watch(event_for_file("baz.py"));
 
         case.apply_changes(changes);
 
@@ -1190,7 +1238,7 @@ mod unix {
         update_file(&patched_bar_baz, "def baz(): print('Version 2')")
             .context("Failed to update bar/baz.py")?;
 
-        let changes = case.stop_watch();
+        let changes = case.stop_watch(event_for_file("baz.py"));
 
         case.apply_changes(changes);
 
@@ -1298,7 +1346,7 @@ mod unix {
         update_file(&baz_original, "def baz(): print('Version 2')")
             .context("Failed to update bar/baz.py")?;
 
-        let changes = case.stop_watch();
+        let changes = case.stop_watch(event_for_file("baz.py"));
 
         case.apply_changes(changes);
 
@@ -1352,7 +1400,7 @@ fn nested_packages_delete_root() -> anyhow::Result<()> {
 
     std::fs::remove_file(case.workspace_path("pyproject.toml").as_std_path())?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(ChangeEvent::is_deleted);
 
     case.apply_changes(changes);
 
@@ -1364,7 +1412,6 @@ fn nested_packages_delete_root() -> anyhow::Result<()> {
 
 #[test]
 fn added_package() -> anyhow::Result<()> {
-    let _ = setup_logging();
     let mut case = setup([
         (
             "pyproject.toml",
@@ -1406,7 +1453,7 @@ fn added_package() -> anyhow::Result<()> {
     )
     .context("failed to write pyproject.toml for package b")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(event_for_file("pyproject.toml"));
 
     case.apply_changes(changes);
 
@@ -1449,7 +1496,7 @@ fn removed_package() -> anyhow::Result<()> {
     std::fs::remove_dir_all(case.workspace_path("packages/b").as_std_path())
         .context("failed to remove package 'b'")?;
 
-    let changes = case.stop_watch();
+    let changes = case.stop_watch(ChangeEvent::is_deleted);
 
     case.apply_changes(changes);
 
diff --git a/crates/red_knot_workspace/src/watch.rs b/crates/red_knot_workspace/src/watch.rs
index a0d1208822..65e0149165 100644
--- a/crates/red_knot_workspace/src/watch.rs
+++ b/crates/red_knot_workspace/src/watch.rs
@@ -81,6 +81,22 @@ impl ChangeEvent {
             _ => None,
         }
     }
+
+    pub const fn is_rescan(&self) -> bool {
+        matches!(self, ChangeEvent::Rescan)
+    }
+
+    pub const fn is_created(&self) -> bool {
+        matches!(self, ChangeEvent::Created { .. })
+    }
+
+    pub const fn is_changed(&self) -> bool {
+        matches!(self, ChangeEvent::Changed { .. })
+    }
+
+    pub const fn is_deleted(&self) -> bool {
+        matches!(self, ChangeEvent::Deleted { .. })
+    }
 }
 
 /// Classification of an event that creates a new path.