From bba8d2dc0255dd18b8ae9c29a15cd02b05485217 Mon Sep 17 00:00:00 2001
From: roife <roifewu@gmail.com>
Date: Thu, 26 Dec 2024 20:08:12 +0800
Subject: [PATCH] move env vars for snapshot tests to UpdateTest

---
 crates/ide/src/runnables.rs                  | 24 +++++++++++++++++---
 crates/rust-analyzer/src/handlers/request.rs | 14 ++++++------
 crates/rust-analyzer/src/lsp/to_proto.rs     | 10 ++++----
 3 files changed, 33 insertions(+), 15 deletions(-)

diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 0a987b8eba..873bde153a 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -1,4 +1,4 @@
-use std::{fmt, ops::Not};
+use std::fmt;
 
 use ast::HasName;
 use cfg::{CfgAtom, CfgExpr};
@@ -20,7 +20,7 @@ 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};
@@ -639,7 +639,25 @@ impl UpdateTest {
         }
 
         let res: SmolStr = builder.join(" + ").into();
-        res.is_empty().not().then_some(res)
+        if res.is_empty() {
+            None
+        } else {
+            Some(format_smolstr!("↺\u{fe0e} Update Tests ({res})"))
+        }
+    }
+
+    pub fn env(&self) -> SmallVec<[(&str, &str); 3]> {
+        let mut env = SmallVec::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
     }
 }
 
diff --git a/crates/rust-analyzer/src/handlers/request.rs b/crates/rust-analyzer/src/handlers/request.rs
index d18bc728b5..d6da807333 100644
--- a/crates/rust-analyzer/src/handlers/request.rs
+++ b/crates/rust-analyzer/src/handlers/request.rs
@@ -943,7 +943,8 @@ pub(crate) fn handle_runnables(
 
         let update_test = runnable.update_test;
         if let Some(mut runnable) = to_proto::runnable(&snap, runnable)? {
-            if let Some(runnable) = to_proto::make_update_runnable(&runnable, &update_test.label())
+            if let Some(runnable) =
+                to_proto::make_update_runnable(&runnable, &update_test.label(), &update_test.env())
             {
                 res.push(runnable);
             }
@@ -2135,10 +2136,7 @@ fn runnable_action_links(
     }
 
     let client_commands_config = snap.config.client_commands();
-    if !(client_commands_config.run_single
-        || client_commands_config.debug_single
-        || client_commands_config.update_single)
-    {
+    if !(client_commands_config.run_single || client_commands_config.debug_single) {
         return None;
     }
 
@@ -2158,8 +2156,10 @@ fn runnable_action_links(
         group.commands.push(to_command_link(dbg_command, r.label.clone()));
     }
 
-    if hover_actions_config.update_test && client_commands_config.update_single {
-        if let Some(update_command) = to_proto::command::update_single(&r, &update_test.label()) {
+    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()));
         }
     }
diff --git a/crates/rust-analyzer/src/lsp/to_proto.rs b/crates/rust-analyzer/src/lsp/to_proto.rs
index 091ac77332..4533755bb4 100644
--- a/crates/rust-analyzer/src/lsp/to_proto.rs
+++ b/crates/rust-analyzer/src/lsp/to_proto.rs
@@ -1606,7 +1606,8 @@ pub(crate) fn code_lens(
                     }
                     if lens_config.update_test && client_commands_config.run_single {
                         let label = update_test.label();
-                        if let Some(r) = make_update_runnable(&r, &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,
@@ -1851,9 +1852,10 @@ pub(crate) mod command {
     }
 }
 
-fn make_update_runnable(
+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;
@@ -1867,9 +1869,7 @@ fn make_update_runnable(
         unreachable!();
     };
 
-    let environment_vars =
-        [("UPDATE_EXPECT", "1"), ("INSTA_UPDATE", "always"), ("SNAPSHOTS", "overwrite")];
-    r.environment.extend(environment_vars.into_iter().map(|(k, v)| (k.to_owned(), v.to_owned())));
+    r.environment.extend(env.iter().map(|(k, v)| (k.to_string(), v.to_string())));
 
     Some(runnable)
 }