Merge pull request #18757 from roife/fix-17812

feat: support updating snapshot tests with codelens/hovering/runnables

13 files changed, 540 insertions, 97 deletions

diff --git a/src/tools/rust-analyzer/crates/hir/src/lib.rs b/src/tools/rust-analyzer/crates/hir/src/lib.rs
index f8af04302f0..f0d8d5a8301 100644
--- a/src/tools/rust-analyzer/crates/hir/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/lib.rs
@@ -5933,6 +5933,12 @@ impl HasCrate for Adt {
     }
 }
 
+impl HasCrate for Impl {
+    fn krate(&self, db: &dyn HirDatabase) -> Crate {
+        self.module(db).krate()
+    }
+}
+
 impl HasCrate for Module {
     fn krate(&self, _: &dyn HirDatabase) -> Crate {
         Module::krate(*self)
diff --git a/src/tools/rust-analyzer/crates/ide/src/annotations.rs b/src/tools/rust-analyzer/crates/ide/src/annotations.rs
index 121a463c9f1..6a4e5ba290e 100644
--- a/src/tools/rust-analyzer/crates/ide/src/annotations.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/annotations.rs
@@ -316,6 +316,11 @@ fn main() {
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -401,6 +406,11 @@ fn main() {
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -537,6 +547,11 @@ fn main() {
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -597,6 +612,11 @@ fn main() {}
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -709,6 +729,11 @@ fn main() {
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -746,6 +771,20 @@ mod tests {
                 [
                     Annotation {
                         range: 3..7,
+                        kind: HasReferences {
+                            pos: FilePositionWrapper {
+                                file_id: FileId(
+                                    0,
+                                ),
+                                offset: 3,
+                            },
+                            data: Some(
+                                [],
+                            ),
+                        },
+                    },
+                    Annotation {
+                        range: 3..7,
                         kind: Runnable(
                             Runnable {
                                 use_name_in_title: false,
@@ -760,24 +799,15 @@ mod tests {
                                 },
                                 kind: Bin,
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
                     Annotation {
-                        range: 3..7,
-                        kind: HasReferences {
-                            pos: FilePositionWrapper {
-                                file_id: FileId(
-                                    0,
-                                ),
-                                offset: 3,
-                            },
-                            data: Some(
-                                [],
-                            ),
-                        },
-                    },
-                    Annotation {
                         range: 18..23,
                         kind: Runnable(
                             Runnable {
@@ -796,6 +826,11 @@ mod tests {
                                     path: "tests",
                                 },
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
@@ -822,6 +857,11 @@ mod tests {
                                     },
                                 },
                                 cfg: None,
+                                update_test: UpdateTest {
+                                    expect_test: false,
+                                    insta: false,
+                                    snapbox: false,
+                                },
                             },
                         ),
                     },
diff --git a/src/tools/rust-analyzer/crates/ide/src/hover/tests.rs b/src/tools/rust-analyzer/crates/ide/src/hover/tests.rs
index aca7bd37511..4154572383e 100644
--- a/src/tools/rust-analyzer/crates/ide/src/hover/tests.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/hover/tests.rs
@@ -3260,6 +3260,11 @@ fn foo_$0test() {}
                             },
                         },
                         cfg: None,
+                        update_test: UpdateTest {
+                            expect_test: false,
+                            insta: false,
+                            snapbox: false,
+                        },
                     },
                 ),
             ]
@@ -3277,28 +3282,33 @@ mod tests$0 {
 }
 "#,
         expect![[r#"
-                [
-                    Runnable(
-                        Runnable {
-                            use_name_in_title: false,
-                            nav: NavigationTarget {
-                                file_id: FileId(
-                                    0,
-                                ),
-                                full_range: 0..46,
-                                focus_range: 4..9,
-                                name: "tests",
-                                kind: Module,
-                                description: "mod tests",
-                            },
-                            kind: TestMod {
-                                path: "tests",
-                            },
-                            cfg: None,
+            [
+                Runnable(
+                    Runnable {
+                        use_name_in_title: false,
+                        nav: NavigationTarget {
+                            file_id: FileId(
+                                0,
+                            ),
+                            full_range: 0..46,
+                            focus_range: 4..9,
+                            name: "tests",
+                            kind: Module,
+                            description: "mod tests",
                         },
-                    ),
-                ]
-            "#]],
+                        kind: TestMod {
+                            path: "tests",
+                        },
+                        cfg: None,
+                        update_test: UpdateTest {
+                            expect_test: false,
+                            insta: false,
+                            snapbox: false,
+                        },
+                    },
+                ),
+            ]
+        "#]],
     );
 }
 
@@ -10029,3 +10039,99 @@ fn bar() {
         "#]],
     );
 }
+
+#[test]
+fn test_runnables_with_snapshot_tests() {
+    check_actions(
+        r#"
+//- /lib.rs crate:foo deps:expect_test,insta,snapbox
+use expect_test::expect;
+use insta::assert_debug_snapshot;
+use snapbox::Assert;
+
+#[test]
+fn test$0() {
+    let actual = "new25";
+    expect!["new25"].assert_eq(&actual);
+    Assert::new()
+        .action_env("SNAPSHOTS")
+        .eq(actual, snapbox::str!["new25"]);
+    assert_debug_snapshot!(actual);
+}
+
+//- /lib.rs crate:expect_test
+struct Expect;
+
+impl Expect {
+    fn assert_eq(&self, actual: &str) {}
+}
+
+#[macro_export]
+macro_rules! expect {
+    ($e:expr) => Expect; // dummy
+}
+
+//- /lib.rs crate:insta
+#[macro_export]
+macro_rules! assert_debug_snapshot {
+    ($e:expr) => {}; // dummy
+}
+
+//- /lib.rs crate:snapbox
+pub struct Assert;
+
+impl Assert {
+    pub fn new() -> Self { Assert }
+
+    pub fn action_env(&self, env: &str) -> &Self { self }
+
+    pub fn eq(&self, actual: &str, expected: &str) {}
+}
+
+#[macro_export]
+macro_rules! str {
+    ($e:expr) => ""; // dummy
+}
+        "#,
+        expect![[r#"
+            [
+                Reference(
+                    FilePositionWrapper {
+                        file_id: FileId(
+                            0,
+                        ),
+                        offset: 92,
+                    },
+                ),
+                Runnable(
+                    Runnable {
+                        use_name_in_title: false,
+                        nav: NavigationTarget {
+                            file_id: FileId(
+                                0,
+                            ),
+                            full_range: 81..301,
+                            focus_range: 92..96,
+                            name: "test",
+                            kind: Function,
+                        },
+                        kind: Test {
+                            test_id: Path(
+                                "test",
+                            ),
+                            attr: TestAttr {
+                                ignore: false,
+                            },
+                        },
+                        cfg: None,
+                        update_test: UpdateTest {
+                            expect_test: true,
+                            insta: true,
+                            snapbox: true,
+                        },
+                    },
+                ),
+            ]
+        "#]],
+    );
+}
diff --git a/src/tools/rust-analyzer/crates/ide/src/runnables.rs b/src/tools/rust-analyzer/crates/ide/src/runnables.rs
index d385e453e21..e89a6339026 100644
--- a/src/tools/rust-analyzer/crates/ide/src/runnables.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/runnables.rs
@@ -1,10 +1,11 @@
-use std::fmt;
+use std::{fmt, sync::OnceLock};
 
+use arrayvec::ArrayVec;
 use ast::HasName;
 use cfg::{CfgAtom, CfgExpr};
 use hir::{
     db::HirDatabase, sym, AsAssocItem, AttrsWithOwner, HasAttrs, HasCrate, HasSource, HirFileIdExt,
-    Semantics,
+    ModPath, Name, PathKind, Semantics, Symbol,
 };
 use ide_assists::utils::{has_test_related_attribute, test_related_attribute_syn};
 use ide_db::{
@@ -15,11 +16,12 @@ use ide_db::{
     FilePosition, FxHashMap, FxHashSet, RootDatabase, SymbolKind,
 };
 use itertools::Itertools;
+use smallvec::SmallVec;
 use span::{Edition, TextSize};
 use stdx::format_to;
 use syntax::{
     ast::{self, AstNode},
-    SmolStr, SyntaxNode, ToSmolStr,
+    format_smolstr, SmolStr, SyntaxNode, ToSmolStr,
 };
 
 use crate::{references, FileId, NavigationTarget, ToNav, TryToNav};
@@ -30,6 +32,7 @@ pub struct Runnable {
     pub nav: NavigationTarget,
     pub kind: RunnableKind,
     pub cfg: Option<CfgExpr>,
+    pub update_test: UpdateTest,
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -334,14 +337,20 @@ pub(crate) fn runnable_fn(
         }
     };
 
+    let fn_source = sema.source(def)?;
     let nav = NavigationTarget::from_named(
         sema.db,
-        def.source(sema.db)?.as_ref().map(|it| it as &dyn ast::HasName),
+        fn_source.as_ref().map(|it| it as &dyn ast::HasName),
         SymbolKind::Function,
     )
     .call_site();
+
+    let file_range = fn_source.syntax().original_file_range_with_macro_call_body(sema.db);
+    let update_test =
+        UpdateTest::find_snapshot_macro(sema, &fn_source.file_syntax(sema.db), file_range);
+
     let cfg = def.attrs(sema.db).cfg();
-    Some(Runnable { use_name_in_title: false, nav, kind, cfg })
+    Some(Runnable { use_name_in_title: false, nav, kind, cfg, update_test })
 }
 
 pub(crate) fn runnable_mod(
@@ -366,7 +375,22 @@ pub(crate) fn runnable_mod(
     let attrs = def.attrs(sema.db);
     let cfg = attrs.cfg();
     let nav = NavigationTarget::from_module_to_decl(sema.db, def).call_site();
-    Some(Runnable { use_name_in_title: false, nav, kind: RunnableKind::TestMod { path }, cfg })
+
+    let module_source = sema.module_definition_node(def);
+    let module_syntax = module_source.file_syntax(sema.db);
+    let file_range = hir::FileRange {
+        file_id: module_source.file_id.original_file(sema.db),
+        range: module_syntax.text_range(),
+    };
+    let update_test = UpdateTest::find_snapshot_macro(sema, &module_syntax, file_range);
+
+    Some(Runnable {
+        use_name_in_title: false,
+        nav,
+        kind: RunnableKind::TestMod { path },
+        cfg,
+        update_test,
+    })
 }
 
 pub(crate) fn runnable_impl(
@@ -392,7 +416,19 @@ pub(crate) fn runnable_impl(
     test_id.retain(|c| c != ' ');
     let test_id = TestId::Path(test_id);
 
-    Some(Runnable { use_name_in_title: false, nav, kind: RunnableKind::DocTest { test_id }, cfg })
+    let impl_source = sema.source(*def)?;
+    let impl_syntax = impl_source.syntax();
+    let file_range = impl_syntax.original_file_range_with_macro_call_body(sema.db);
+    let update_test =
+        UpdateTest::find_snapshot_macro(sema, &impl_syntax.file_syntax(sema.db), file_range);
+
+    Some(Runnable {
+        use_name_in_title: false,
+        nav,
+        kind: RunnableKind::DocTest { test_id },
+        cfg,
+        update_test,
+    })
 }
 
 fn has_cfg_test(attrs: AttrsWithOwner) -> bool {
@@ -404,6 +440,8 @@ fn runnable_mod_outline_definition(
     sema: &Semantics<'_, RootDatabase>,
     def: hir::Module,
 ) -> Option<Runnable> {
+    def.as_source_file_id(sema.db)?;
+
     if !has_test_function_or_multiple_test_submodules(sema, &def, has_cfg_test(def.attrs(sema.db)))
     {
         return None;
@@ -421,16 +459,22 @@ fn runnable_mod_outline_definition(
 
     let attrs = def.attrs(sema.db);
     let cfg = attrs.cfg();
-    if def.as_source_file_id(sema.db).is_some() {
-        Some(Runnable {
-            use_name_in_title: false,
-            nav: def.to_nav(sema.db).call_site(),
-            kind: RunnableKind::TestMod { path },
-            cfg,
-        })
-    } else {
-        None
-    }
+
+    let mod_source = sema.module_definition_node(def);
+    let mod_syntax = mod_source.file_syntax(sema.db);
+    let file_range = hir::FileRange {
+        file_id: mod_source.file_id.original_file(sema.db),
+        range: mod_syntax.text_range(),
+    };
+    let update_test = UpdateTest::find_snapshot_macro(sema, &mod_syntax, file_range);
+
+    Some(Runnable {
+        use_name_in_title: false,
+        nav: def.to_nav(sema.db).call_site(),
+        kind: RunnableKind::TestMod { path },
+        cfg,
+        update_test,
+    })
 }
 
 fn module_def_doctest(db: &RootDatabase, def: Definition) -> Option<Runnable> {
@@ -495,6 +539,7 @@ fn module_def_doctest(db: &RootDatabase, def: Definition) -> Option<Runnable> {
         nav,
         kind: RunnableKind::DocTest { test_id },
         cfg: attrs.cfg(),
+        update_test: UpdateTest::default(),
     };
     Some(res)
 }
@@ -575,6 +620,128 @@ fn has_test_function_or_multiple_test_submodules(
     number_of_test_submodules > 1
 }
 
+#[derive(Debug, Default, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct UpdateTest {
+    pub expect_test: bool,
+    pub insta: bool,
+    pub snapbox: bool,
+}
+
+static SNAPSHOT_TEST_MACROS: OnceLock<FxHashMap<&str, Vec<ModPath>>> = OnceLock::new();
+
+impl UpdateTest {
+    const EXPECT_CRATE: &str = "expect_test";
+    const EXPECT_MACROS: &[&str] = &["expect", "expect_file"];
+
+    const INSTA_CRATE: &str = "insta";
+    const INSTA_MACROS: &[&str] = &[
+        "assert_snapshot",
+        "assert_debug_snapshot",
+        "assert_display_snapshot",
+        "assert_json_snapshot",
+        "assert_yaml_snapshot",
+        "assert_ron_snapshot",
+        "assert_toml_snapshot",
+        "assert_csv_snapshot",
+        "assert_compact_json_snapshot",
+        "assert_compact_debug_snapshot",
+        "assert_binary_snapshot",
+    ];
+
+    const SNAPBOX_CRATE: &str = "snapbox";
+    const SNAPBOX_MACROS: &[&str] = &["assert_data_eq", "file", "str"];
+
+    fn find_snapshot_macro(
+        sema: &Semantics<'_, RootDatabase>,
+        scope: &SyntaxNode,
+        file_range: hir::FileRange,
+    ) -> Self {
+        fn init<'a>(
+            krate_name: &'a str,
+            paths: &[&str],
+            map: &mut FxHashMap<&'a str, Vec<ModPath>>,
+        ) {
+            let mut res = Vec::with_capacity(paths.len());
+            let krate = Name::new_symbol_root(Symbol::intern(krate_name));
+            for path in paths {
+                let segments = [krate.clone(), Name::new_symbol_root(Symbol::intern(path))];
+                let mod_path = ModPath::from_segments(PathKind::Abs, segments);
+                res.push(mod_path);
+            }
+            map.insert(krate_name, res);
+        }
+
+        let mod_paths = SNAPSHOT_TEST_MACROS.get_or_init(|| {
+            let mut map = FxHashMap::default();
+            init(Self::EXPECT_CRATE, Self::EXPECT_MACROS, &mut map);
+            init(Self::INSTA_CRATE, Self::INSTA_MACROS, &mut map);
+            init(Self::SNAPBOX_CRATE, Self::SNAPBOX_MACROS, &mut map);
+            map
+        });
+
+        let search_scope = SearchScope::file_range(file_range);
+        let find_macro = |paths: &[ModPath]| {
+            for path in paths {
+                let Some(items) = sema.resolve_mod_path(scope, path) else {
+                    continue;
+                };
+                for item in items {
+                    if let hir::ItemInNs::Macros(makro) = item {
+                        if Definition::Macro(makro)
+                            .usages(sema)
+                            .in_scope(&search_scope)
+                            .at_least_one()
+                        {
+                            return true;
+                        }
+                    }
+                }
+            }
+            false
+        };
+
+        UpdateTest {
+            expect_test: find_macro(mod_paths.get(Self::EXPECT_CRATE).unwrap()),
+            insta: find_macro(mod_paths.get(Self::INSTA_CRATE).unwrap()),
+            snapbox: find_macro(mod_paths.get(Self::SNAPBOX_CRATE).unwrap()),
+        }
+    }
+
+    pub fn label(&self) -> Option<SmolStr> {
+        let mut builder: SmallVec<[_; 3]> = SmallVec::new();
+        if self.expect_test {
+            builder.push("Expect");
+        }
+        if self.insta {
+            builder.push("Insta");
+        }
+        if self.snapbox {
+            builder.push("Snapbox");
+        }
+
+        let res: SmolStr = builder.join(" + ").into();
+        if res.is_empty() {
+            None
+        } else {
+            Some(format_smolstr!("↺\u{fe0e} Update Tests ({res})"))
+        }
+    }
+
+    pub fn env(&self) -> ArrayVec<(&str, &str), 3> {
+        let mut env = ArrayVec::new();
+        if self.expect_test {
+            env.push(("UPDATE_EXPECT", "1"));
+        }
+        if self.insta {
+            env.push(("INSTA_UPDATE", "always"));
+        }
+        if self.snapbox {
+            env.push(("SNAPSHOTS", "overwrite"));
+        }
+        env
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use expect_test::{expect, Expect};
@@ -1337,18 +1504,18 @@ mod tests {
                         file_id: FileId(
                             0,
                         ),
-                        full_range: 52..115,
-                        focus_range: 67..75,
-                        name: "foo_test",
+                        full_range: 121..185,
+                        focus_range: 136..145,
+                        name: "foo2_test",
                         kind: Function,
                     },
                     NavigationTarget {
                         file_id: FileId(
                             0,
                         ),
-                        full_range: 121..185,
-                        focus_range: 136..145,
-                        name: "foo2_test",
+                        full_range: 52..115,
+                        focus_range: 67..75,
+                        name: "foo_test",
                         kind: Function,
                     },
                 ]
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/config.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/config.rs
index c182952c731..0f8840a810c 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/config.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/config.rs
@@ -119,6 +119,9 @@ config_data! {
         /// Whether to show `Run` action. Only applies when
         /// `#rust-analyzer.hover.actions.enable#` is set.
         hover_actions_run_enable: bool             = true,
+        /// Whether to show `Update Test` action. Only applies when
+        /// `#rust-analyzer.hover.actions.enable#` and `#rust-analyzer.hover.actions.run.enable#` are set.
+        hover_actions_updateTest_enable: bool     = true,
 
         /// Whether to show documentation on hover.
         hover_documentation_enable: bool           = true,
@@ -243,6 +246,9 @@ config_data! {
         /// Whether to show `Run` lens. Only applies when
         /// `#rust-analyzer.lens.enable#` is set.
         lens_run_enable: bool              = true,
+        /// Whether to show `Update Test` lens. Only applies when
+        /// `#rust-analyzer.lens.enable#` and `#rust-analyzer.lens.run.enable#` are set.
+        lens_updateTest_enable: bool = true,
 
         /// Disable project auto-discovery in favor of explicitly specified set
         /// of projects.
@@ -1161,6 +1167,7 @@ pub struct LensConfig {
     // runnables
     pub run: bool,
     pub debug: bool,
+    pub update_test: bool,
     pub interpret: bool,
 
     // implementations
@@ -1196,6 +1203,7 @@ impl LensConfig {
     pub fn any(&self) -> bool {
         self.run
             || self.debug
+            || self.update_test
             || self.implementations
             || self.method_refs
             || self.refs_adt
@@ -1208,7 +1216,7 @@ impl LensConfig {
     }
 
     pub fn runnable(&self) -> bool {
-        self.run || self.debug
+        self.run || self.debug || self.update_test
     }
 
     pub fn references(&self) -> bool {
@@ -1222,6 +1230,7 @@ pub struct HoverActionsConfig {
     pub references: bool,
     pub run: bool,
     pub debug: bool,
+    pub update_test: bool,
     pub goto_type_def: bool,
 }
 
@@ -1231,6 +1240,7 @@ impl HoverActionsConfig {
         references: false,
         run: false,
         debug: false,
+        update_test: false,
         goto_type_def: false,
     };
 
@@ -1243,7 +1253,7 @@ impl HoverActionsConfig {
     }
 
     pub fn runnable(&self) -> bool {
-        self.run || self.debug
+        self.run || self.debug || self.update_test
     }
 }
 
@@ -1517,6 +1527,9 @@ impl Config {
             references: enable && self.hover_actions_references_enable().to_owned(),
             run: enable && self.hover_actions_run_enable().to_owned(),
             debug: enable && self.hover_actions_debug_enable().to_owned(),
+            update_test: enable
+                && self.hover_actions_run_enable().to_owned()
+                && self.hover_actions_updateTest_enable().to_owned(),
             goto_type_def: enable && self.hover_actions_gotoTypeDef_enable().to_owned(),
         }
     }
@@ -2120,6 +2133,9 @@ impl Config {
         LensConfig {
             run: *self.lens_enable() && *self.lens_run_enable(),
             debug: *self.lens_enable() && *self.lens_debug_enable(),
+            update_test: *self.lens_enable()
+                && *self.lens_updateTest_enable()
+                && *self.lens_run_enable(),
             interpret: *self.lens_enable() && *self.lens_run_enable() && *self.interpret_tests(),
             implementations: *self.lens_enable() && *self.lens_implementations_enable(),
             method_refs: *self.lens_enable() && *self.lens_references_method_enable(),
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs
index f103f6cbe27..7ac70efe2d6 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/handlers/request.rs
@@ -27,7 +27,7 @@ use paths::Utf8PathBuf;
 use project_model::{CargoWorkspace, ManifestPath, ProjectWorkspaceKind, TargetKind};
 use serde_json::json;
 use stdx::{format_to, never};
-use syntax::{algo, ast, AstNode, TextRange, TextSize};
+use syntax::{TextRange, TextSize};
 use triomphe::Arc;
 use vfs::{AbsPath, AbsPathBuf, FileId, VfsPath};
 
@@ -928,39 +928,32 @@ pub(crate) fn handle_runnables(
     let offset = params.position.and_then(|it| from_proto::offset(&line_index, it).ok());
     let target_spec = TargetSpec::for_file(&snap, file_id)?;
 
-    let expect_test = match offset {
-        Some(offset) => {
-            let source_file = snap.analysis.parse(file_id)?;
-            algo::find_node_at_offset::<ast::MacroCall>(source_file.syntax(), offset)
-                .and_then(|it| it.path()?.segment()?.name_ref())
-                .map_or(false, |it| it.text() == "expect" || it.text() == "expect_file")
-        }
-        None => false,
-    };
-
     let mut res = Vec::new();
     for runnable in snap.analysis.runnables(file_id)? {
-        if should_skip_for_offset(&runnable, offset) {
-            continue;
-        }
-        if should_skip_target(&runnable, target_spec.as_ref()) {
+        if should_skip_for_offset(&runnable, offset)
+            || should_skip_target(&runnable, target_spec.as_ref())
+        {
             continue;
         }
+
+        let update_test = runnable.update_test;
         if let Some(mut runnable) = to_proto::runnable(&snap, runnable)? {
-            if expect_test {
-                if let lsp_ext::RunnableArgs::Cargo(r) = &mut runnable.args {
-                    runnable.label = format!("{} + expect", runnable.label);
-                    r.environment.insert("UPDATE_EXPECT".to_owned(), "1".to_owned());
-                    if let Some(TargetSpec::Cargo(CargoTargetSpec {
-                        sysroot_root: Some(sysroot_root),
-                        ..
-                    })) = &target_spec
-                    {
-                        r.environment
-                            .insert("RUSTC_TOOLCHAIN".to_owned(), sysroot_root.to_string());
-                    }
-                }
+            if let Some(runnable) =
+                to_proto::make_update_runnable(&runnable, &update_test.label(), &update_test.env())
+            {
+                res.push(runnable);
             }
+
+            if let lsp_ext::RunnableArgs::Cargo(r) = &mut runnable.args {
+                if let Some(TargetSpec::Cargo(CargoTargetSpec {
+                    sysroot_root: Some(sysroot_root),
+                    ..
+                })) = &target_spec
+                {
+                    r.environment.insert("RUSTC_TOOLCHAIN".to_owned(), sysroot_root.to_string());
+                }
+            };
+
             res.push(runnable);
         }
     }
@@ -2142,6 +2135,7 @@ fn runnable_action_links(
     }
 
     let title = runnable.title();
+    let update_test = runnable.update_test;
     let r = to_proto::runnable(snap, runnable).ok()??;
 
     let mut group = lsp_ext::CommandLinkGroup::default();
@@ -2153,7 +2147,15 @@ fn runnable_action_links(
 
     if hover_actions_config.debug && client_commands_config.debug_single {
         let dbg_command = to_proto::command::debug_single(&r);
-        group.commands.push(to_command_link(dbg_command, r.label));
+        group.commands.push(to_command_link(dbg_command, r.label.clone()));
+    }
+
+    if hover_actions_config.update_test && client_commands_config.run_single {
+        let label = update_test.label();
+        if let Some(r) = to_proto::make_update_runnable(&r, &label, &update_test.env()) {
+            let update_command = to_proto::command::run_single(&r, label.unwrap().as_str());
+            group.commands.push(to_command_link(update_command, r.label.clone()));
+        }
     }
 
     Some(group)
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/ext.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/ext.rs
index c0173d9c247..e1677cbcda8 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/ext.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/ext.rs
@@ -427,14 +427,14 @@ impl Request for Runnables {
     const METHOD: &'static str = "experimental/runnables";
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct RunnablesParams {
     pub text_document: TextDocumentIdentifier,
     pub position: Option<Position>,
 }
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct Runnable {
     pub label: String,
@@ -444,7 +444,7 @@ pub struct Runnable {
     pub args: RunnableArgs,
 }
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 #[serde(untagged)]
 pub enum RunnableArgs {
@@ -452,14 +452,14 @@ pub enum RunnableArgs {
     Shell(ShellRunnableArgs),
 }
 
-#[derive(Serialize, Deserialize, Debug)]
+#[derive(Serialize, Deserialize, Debug, Clone)]
 #[serde(rename_all = "lowercase")]
 pub enum RunnableKind {
     Cargo,
     Shell,
 }
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct CargoRunnableArgs {
     #[serde(skip_serializing_if = "FxHashMap::is_empty")]
@@ -475,7 +475,7 @@ pub struct CargoRunnableArgs {
     pub executable_args: Vec<String>,
 }
 
-#[derive(Deserialize, Serialize, Debug)]
+#[derive(Deserialize, Serialize, Debug, Clone)]
 #[serde(rename_all = "camelCase")]
 pub struct ShellRunnableArgs {
     #[serde(skip_serializing_if = "FxHashMap::is_empty")]
diff --git a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
index 05e93b4e6ac..4533755bb4f 100644
--- a/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
+++ b/src/tools/rust-analyzer/crates/rust-analyzer/src/lsp/to_proto.rs
@@ -20,6 +20,7 @@ use itertools::Itertools;
 use paths::{Utf8Component, Utf8Prefix};
 use semver::VersionReq;
 use serde_json::to_value;
+use syntax::SmolStr;
 use vfs::AbsPath;
 
 use crate::{
@@ -1567,6 +1568,7 @@ pub(crate) fn code_lens(
             let line_index = snap.file_line_index(run.nav.file_id)?;
             let annotation_range = range(&line_index, annotation.range);
 
+            let update_test = run.update_test;
             let title = run.title();
             let can_debug = match run.kind {
                 ide::RunnableKind::DocTest { .. } => false,
@@ -1602,6 +1604,18 @@ pub(crate) fn code_lens(
                             data: None,
                         })
                     }
+                    if lens_config.update_test && client_commands_config.run_single {
+                        let label = update_test.label();
+                        let env = update_test.env();
+                        if let Some(r) = make_update_runnable(&r, &label, &env) {
+                            let command = command::run_single(&r, label.unwrap().as_str());
+                            acc.push(lsp_types::CodeLens {
+                                range: annotation_range,
+                                command: Some(command),
+                                data: None,
+                            })
+                        }
+                    }
                 }
 
                 if lens_config.interpret {
@@ -1786,7 +1800,7 @@ pub(crate) mod command {
 
     pub(crate) fn debug_single(runnable: &lsp_ext::Runnable) -> lsp_types::Command {
         lsp_types::Command {
-            title: "Debug".into(),
+            title: "⚙\u{fe0e} Debug".into(),
             command: "rust-analyzer.debugSingle".into(),
             arguments: Some(vec![to_value(runnable).unwrap()]),
         }
@@ -1838,6 +1852,28 @@ pub(crate) mod command {
     }
 }
 
+pub(crate) fn make_update_runnable(
+    runnable: &lsp_ext::Runnable,
+    label: &Option<SmolStr>,
+    env: &[(&str, &str)],
+) -> Option<lsp_ext::Runnable> {
+    if !matches!(runnable.args, lsp_ext::RunnableArgs::Cargo(_)) {
+        return None;
+    }
+    let label = label.as_ref()?;
+
+    let mut runnable = runnable.clone();
+    runnable.label = format!("{} + {}", runnable.label, label);
+
+    let lsp_ext::RunnableArgs::Cargo(r) = &mut runnable.args else {
+        unreachable!();
+    };
+
+    r.environment.extend(env.iter().map(|(k, v)| (k.to_string(), v.to_string())));
+
+    Some(runnable)
+}
+
 pub(crate) fn implementation_title(count: usize) -> String {
     if count == 1 {
         "1 implementation".into()
diff --git a/src/tools/rust-analyzer/docs/dev/lsp-extensions.md b/src/tools/rust-analyzer/docs/dev/lsp-extensions.md
index 0e37611a549..826ce112448 100644
--- a/src/tools/rust-analyzer/docs/dev/lsp-extensions.md
+++ b/src/tools/rust-analyzer/docs/dev/lsp-extensions.md
@@ -1,5 +1,5 @@
 <!---
-lsp/ext.rs hash: 9790509d87670c22
+lsp/ext.rs hash: 512c06cd8b46a21d
 
 If you need to change the above hash to make the test pass, please check if you
 need to adjust this doc as well and ping this issue:
diff --git a/src/tools/rust-analyzer/docs/user/generated_config.adoc b/src/tools/rust-analyzer/docs/user/generated_config.adoc
index 5056c7d977c..fffbb94efab 100644
--- a/src/tools/rust-analyzer/docs/user/generated_config.adoc
+++ b/src/tools/rust-analyzer/docs/user/generated_config.adoc
@@ -497,6 +497,12 @@ Whether to show `References` action. Only applies when
 Whether to show `Run` action. Only applies when
 `#rust-analyzer.hover.actions.enable#` is set.
 --
+[[rust-analyzer.hover.actions.updateTest.enable]]rust-analyzer.hover.actions.updateTest.enable (default: `true`)::
++
+--
+Whether to show `Update Test` action. Only applies when
+`#rust-analyzer.hover.actions.enable#` and `#rust-analyzer.hover.actions.run.enable#` are set.
+--
 [[rust-analyzer.hover.documentation.enable]]rust-analyzer.hover.documentation.enable (default: `true`)::
 +
 --
@@ -808,6 +814,12 @@ Only applies when `#rust-analyzer.lens.enable#` is set.
 Whether to show `Run` lens. Only applies when
 `#rust-analyzer.lens.enable#` is set.
 --
+[[rust-analyzer.lens.updateTest.enable]]rust-analyzer.lens.updateTest.enable (default: `true`)::
++
+--
+Whether to show `Update Test` lens. Only applies when
+`#rust-analyzer.lens.enable#` and `#rust-analyzer.lens.run.enable#` are set.
+--
 [[rust-analyzer.linkedProjects]]rust-analyzer.linkedProjects (default: `[]`)::
 +
 --
diff --git a/src/tools/rust-analyzer/editors/code/package.json b/src/tools/rust-analyzer/editors/code/package.json
index b9249e9ac8b..426414becb3 100644
--- a/src/tools/rust-analyzer/editors/code/package.json
+++ b/src/tools/rust-analyzer/editors/code/package.json
@@ -407,6 +407,11 @@
                             "$rustc"
                         ],
                         "markdownDescription": "Problem matchers to use for `rust-analyzer.run` command, eg `[\"$rustc\", \"$rust-panic\"]`."
+                    },
+                    "rust-analyzer.runnables.askBeforeUpdateTest": {
+                        "type": "boolean",
+                        "default": true,
+                        "markdownDescription": "Ask before updating the test when running it."
                     }
                 }
             },
@@ -1518,6 +1523,16 @@
             {
                 "title": "hover",
                 "properties": {
+                    "rust-analyzer.hover.actions.updateTest.enable": {
+                        "markdownDescription": "Whether to show `Update Test` action. Only applies when\n`#rust-analyzer.hover.actions.enable#` and `#rust-analyzer.hover.actions.run.enable#` are set.",
+                        "default": true,
+                        "type": "boolean"
+                    }
+                }
+            },
+            {
+                "title": "hover",
+                "properties": {
                     "rust-analyzer.hover.documentation.enable": {
                         "markdownDescription": "Whether to show documentation on hover.",
                         "default": true,
@@ -2296,6 +2311,16 @@
                 }
             },
             {
+                "title": "lens",
+                "properties": {
+                    "rust-analyzer.lens.updateTest.enable": {
+                        "markdownDescription": "Whether to show `Update Test` lens. Only applies when\n`#rust-analyzer.lens.enable#` and `#rust-analyzer.lens.run.enable#` are set.",
+                        "default": true,
+                        "type": "boolean"
+                    }
+                }
+            },
+            {
                 "title": "general",
                 "properties": {
                     "rust-analyzer.linkedProjects": {
diff --git a/src/tools/rust-analyzer/editors/code/src/commands.ts b/src/tools/rust-analyzer/editors/code/src/commands.ts
index 7ebc186a3ea..73e39c900e7 100644
--- a/src/tools/rust-analyzer/editors/code/src/commands.ts
+++ b/src/tools/rust-analyzer/editors/code/src/commands.ts
@@ -1139,11 +1139,37 @@ export function peekTests(ctx: CtxInit): Cmd {
     };
 }
 
+function isUpdatingTest(runnable: ra.Runnable): boolean {
+    if (!isCargoRunnableArgs(runnable.args)) {
+        return false;
+    }
+
+    const env = runnable.args.environment;
+    return env ? ["UPDATE_EXPECT", "INSTA_UPDATE", "SNAPSHOTS"].some((key) => key in env) : false;
+}
+
 export function runSingle(ctx: CtxInit): Cmd {
     return async (runnable: ra.Runnable) => {
         const editor = ctx.activeRustEditor;
         if (!editor) return;
 
+        if (isUpdatingTest(runnable) && ctx.config.askBeforeUpdateTest) {
+            const selection = await vscode.window.showInformationMessage(
+                "rust-analyzer",
+                { detail: "Do you want to update tests?", modal: true },
+                "Update Now",
+                "Update (and Don't ask again)",
+            );
+
+            if (selection !== "Update Now" && selection !== "Update (and Don't ask again)") {
+                return;
+            }
+
+            if (selection === "Update (and Don't ask again)") {
+                await ctx.config.setAskBeforeUpdateTest(false);
+            }
+        }
+
         const task = await createTaskFromRunnable(runnable, ctx.config);
         task.group = vscode.TaskGroup.Build;
         task.presentationOptions = {
diff --git a/src/tools/rust-analyzer/editors/code/src/config.ts b/src/tools/rust-analyzer/editors/code/src/config.ts
index a97d4beab51..720c473c5b4 100644
--- a/src/tools/rust-analyzer/editors/code/src/config.ts
+++ b/src/tools/rust-analyzer/editors/code/src/config.ts
@@ -362,6 +362,13 @@ export class Config {
     get initializeStopped() {
         return this.get<boolean>("initializeStopped");
     }
+
+    get askBeforeUpdateTest() {
+        return this.get<boolean>("runnables.askBeforeUpdateTest");
+    }
+    async setAskBeforeUpdateTest(value: boolean) {
+        await this.cfg.update("runnables.askBeforeUpdateTest", value, true);
+    }
 }
 
 export function prepareVSCodeConfig<T>(resp: T): T {