Auto merge of #17014 - Veykril:runnables-exported-main-test, r=Veykril

fix: Consider `exported_name="main"` functions in test modules as tests

Fixes https://github.com/rust-lang/rust-analyzer/issues/17011

6 files changed, 94 insertions, 22 deletions

diff --git a/crates/hir-def/src/attr.rs b/crates/hir-def/src/attr.rs
index fa7730f302e..f4f78ce6a4d 100644
--- a/crates/hir-def/src/attr.rs
+++ b/crates/hir-def/src/attr.rs
@@ -141,6 +141,10 @@ impl Attrs {
         }
     }
 
+    pub fn cfgs(&self) -> impl Iterator<Item = CfgExpr> + '_ {
+        self.by_key("cfg").tt_values().map(CfgExpr::parse)
+    }
+
     pub(crate) fn is_cfg_enabled(&self, cfg_options: &CfgOptions) -> bool {
         match self.cfg() {
             None => true,
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 503c5493274..9b9f710d9a1 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2006,12 +2006,15 @@ impl Function {
 
     /// is this a `fn main` or a function with an `export_name` of `main`?
     pub fn is_main(self, db: &dyn HirDatabase) -> bool {
-        if !self.module(db).is_crate_root() {
-            return false;
-        }
         let data = db.function_data(self.id);
+        data.attrs.export_name() == Some("main")
+            || self.module(db).is_crate_root() && data.name.to_smol_str() == "main"
+    }
 
-        data.name.to_smol_str() == "main" || data.attrs.export_name() == Some("main")
+    /// Is this a function with an `export_name` of `main`?
+    pub fn exported_main(self, db: &dyn HirDatabase) -> bool {
+        let data = db.function_data(self.id);
+        data.attrs.export_name() == Some("main")
     }
 
     /// Does this function have the ignore attribute?
diff --git a/crates/ide-assists/src/handlers/toggle_ignore.rs b/crates/ide-assists/src/handlers/toggle_ignore.rs
index f864ee50c81..264a2f0326e 100644
--- a/crates/ide-assists/src/handlers/toggle_ignore.rs
+++ b/crates/ide-assists/src/handlers/toggle_ignore.rs
@@ -3,7 +3,7 @@ use syntax::{
     AstNode, AstToken,
 };
 
-use crate::{utils::test_related_attribute, AssistContext, AssistId, AssistKind, Assists};
+use crate::{utils::test_related_attribute_syn, AssistContext, AssistId, AssistKind, Assists};
 
 // Assist: toggle_ignore
 //
@@ -26,7 +26,7 @@ use crate::{utils::test_related_attribute, AssistContext, AssistId, AssistKind,
 pub(crate) fn toggle_ignore(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
     let attr: ast::Attr = ctx.find_node_at_offset()?;
     let func = attr.syntax().parent().and_then(ast::Fn::cast)?;
-    let attr = test_related_attribute(&func)?;
+    let attr = test_related_attribute_syn(&func)?;
 
     match has_ignore_attribute(&func) {
         None => acc.add(
diff --git a/crates/ide-assists/src/utils.rs b/crates/ide-assists/src/utils.rs
index 8bd5d179331..bc0c9b79c75 100644
--- a/crates/ide-assists/src/utils.rs
+++ b/crates/ide-assists/src/utils.rs
@@ -71,7 +71,7 @@ pub fn extract_trivial_expression(block_expr: &ast::BlockExpr) -> Option<ast::Ex
 ///
 /// It may produce false positives, for example, `#[wasm_bindgen_test]` requires a different command to run the test,
 /// but it's better than not to have the runnables for the tests at all.
-pub fn test_related_attribute(fn_def: &ast::Fn) -> Option<ast::Attr> {
+pub fn test_related_attribute_syn(fn_def: &ast::Fn) -> Option<ast::Attr> {
     fn_def.attrs().find_map(|attr| {
         let path = attr.path()?;
         let text = path.syntax().text().to_string();
@@ -83,6 +83,19 @@ pub fn test_related_attribute(fn_def: &ast::Fn) -> Option<ast::Attr> {
     })
 }
 
+pub fn has_test_related_attribute(attrs: &hir::AttrsWithOwner) -> bool {
+    attrs.iter().any(|attr| {
+        let path = attr.path();
+        (|| {
+            Some(
+                path.segments().first()?.as_text()?.starts_with("test")
+                    || path.segments().last()?.as_text()?.ends_with("test"),
+            )
+        })()
+        .unwrap_or_default()
+    })
+}
+
 #[derive(Clone, Copy, PartialEq)]
 pub enum IgnoreAssocItems {
     DocHiddenAttrPresent,
diff --git a/crates/ide/src/annotations/fn_references.rs b/crates/ide/src/annotations/fn_references.rs
index a090b60413e..2e7e230e5ac 100644
--- a/crates/ide/src/annotations/fn_references.rs
+++ b/crates/ide/src/annotations/fn_references.rs
@@ -2,7 +2,7 @@
 //! We have to skip tests, so cannot reuse file_structure module.
 
 use hir::Semantics;
-use ide_assists::utils::test_related_attribute;
+use ide_assists::utils::test_related_attribute_syn;
 use ide_db::RootDatabase;
 use syntax::{ast, ast::HasName, AstNode, SyntaxNode, TextRange};
 
@@ -19,7 +19,7 @@ pub(super) fn find_all_methods(
 
 fn method_range(item: SyntaxNode) -> Option<(TextRange, Option<TextRange>)> {
     ast::Fn::cast(item).and_then(|fn_def| {
-        if test_related_attribute(&fn_def).is_some() {
+        if test_related_attribute_syn(&fn_def).is_some() {
             None
         } else {
             Some((
diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 79324bf3877..b6c6753755c 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -1,9 +1,11 @@
 use std::fmt;
 
 use ast::HasName;
-use cfg::CfgExpr;
-use hir::{db::HirDatabase, AsAssocItem, HasAttrs, HasSource, HirFileIdExt, Semantics};
-use ide_assists::utils::test_related_attribute;
+use cfg::{CfgAtom, CfgExpr};
+use hir::{
+    db::HirDatabase, AsAssocItem, AttrsWithOwner, HasAttrs, HasSource, HirFileIdExt, Semantics,
+};
+use ide_assists::utils::{has_test_related_attribute, test_related_attribute_syn};
 use ide_db::{
     base_db::{FilePosition, FileRange},
     defs::Definition,
@@ -280,7 +282,7 @@ fn find_related_tests_in_module(
 }
 
 fn as_test_runnable(sema: &Semantics<'_, RootDatabase>, fn_def: &ast::Fn) -> Option<Runnable> {
-    if test_related_attribute(fn_def).is_some() {
+    if test_related_attribute_syn(fn_def).is_some() {
         let function = sema.to_def(fn_def)?;
         runnable_fn(sema, function)
     } else {
@@ -293,7 +295,7 @@ fn parent_test_module(sema: &Semantics<'_, RootDatabase>, fn_def: &ast::Fn) -> O
         let module = ast::Module::cast(node)?;
         let module = sema.to_def(&module)?;
 
-        if has_test_function_or_multiple_test_submodules(sema, &module) {
+        if has_test_function_or_multiple_test_submodules(sema, &module, false) {
             Some(module)
         } else {
             None
@@ -305,7 +307,8 @@ pub(crate) fn runnable_fn(
     sema: &Semantics<'_, RootDatabase>,
     def: hir::Function,
 ) -> Option<Runnable> {
-    let kind = if def.is_main(sema.db) {
+    let under_cfg_test = has_cfg_test(def.module(sema.db).attrs(sema.db));
+    let kind = if !under_cfg_test && def.is_main(sema.db) {
         RunnableKind::Bin
     } else {
         let test_id = || {
@@ -342,7 +345,8 @@ pub(crate) fn runnable_mod(
     sema: &Semantics<'_, RootDatabase>,
     def: hir::Module,
 ) -> Option<Runnable> {
-    if !has_test_function_or_multiple_test_submodules(sema, &def) {
+    if !has_test_function_or_multiple_test_submodules(sema, &def, has_cfg_test(def.attrs(sema.db)))
+    {
         return None;
     }
     let path = def
@@ -384,12 +388,17 @@ pub(crate) fn runnable_impl(
     Some(Runnable { use_name_in_title: false, nav, kind: RunnableKind::DocTest { test_id }, cfg })
 }
 
+fn has_cfg_test(attrs: AttrsWithOwner) -> bool {
+    attrs.cfgs().any(|cfg| matches!(cfg, CfgExpr::Atom(CfgAtom::Flag(s)) if s == "test"))
+}
+
 /// Creates a test mod runnable for outline modules at the top of their definition.
 fn runnable_mod_outline_definition(
     sema: &Semantics<'_, RootDatabase>,
     def: hir::Module,
 ) -> Option<Runnable> {
-    if !has_test_function_or_multiple_test_submodules(sema, &def) {
+    if !has_test_function_or_multiple_test_submodules(sema, &def, has_cfg_test(def.attrs(sema.db)))
+    {
         return None;
     }
     let path = def
@@ -522,20 +531,28 @@ fn has_runnable_doc_test(attrs: &hir::Attrs) -> bool {
 fn has_test_function_or_multiple_test_submodules(
     sema: &Semantics<'_, RootDatabase>,
     module: &hir::Module,
+    consider_exported_main: bool,
 ) -> bool {
     let mut number_of_test_submodules = 0;
 
     for item in module.declarations(sema.db) {
         match item {
             hir::ModuleDef::Function(f) => {
-                if let Some(it) = f.source(sema.db) {
-                    if test_related_attribute(&it.value).is_some() {
-                        return true;
-                    }
+                if has_test_related_attribute(&f.attrs(sema.db)) {
+                    return true;
+                }
+                if consider_exported_main && f.exported_main(sema.db) {
+                    // an exported main in a test module can be considered a test wrt to custom test
+                    // runners
+                    return true;
                 }
             }
             hir::ModuleDef::Module(submodule) => {
-                if has_test_function_or_multiple_test_submodules(sema, &submodule) {
+                if has_test_function_or_multiple_test_submodules(
+                    sema,
+                    &submodule,
+                    consider_exported_main,
+                ) {
                     number_of_test_submodules += 1;
                 }
             }
@@ -1484,4 +1501,39 @@ mod r#mod {
             "#]],
         )
     }
+
+    #[test]
+    fn exported_main_is_test_in_cfg_test_mod() {
+        check(
+            r#"
+//- /lib.rs crate:foo cfg:test
+$0
+mod not_a_test_module_inline {
+    #[export_name = "main"]
+    fn exp_main() {}
+}
+#[cfg(test)]
+mod test_mod_inline {
+    #[export_name = "main"]
+    fn exp_main() {}
+}
+mod not_a_test_module;
+#[cfg(test)]
+mod test_mod;
+//- /not_a_test_module.rs
+#[export_name = "main"]
+fn exp_main() {}
+//- /test_mod.rs
+#[export_name = "main"]
+fn exp_main() {}
+"#,
+            expect![[r#"
+                [
+                    "(Bin, NavigationTarget { file_id: FileId(0), full_range: 36..80, focus_range: 67..75, name: \"exp_main\", kind: Function })",
+                    "(TestMod, NavigationTarget { file_id: FileId(0), full_range: 83..168, focus_range: 100..115, name: \"test_mod_inline\", kind: Module, description: \"mod test_mod_inline\" }, Atom(Flag(\"test\")))",
+                    "(TestMod, NavigationTarget { file_id: FileId(0), full_range: 192..218, focus_range: 209..217, name: \"test_mod\", kind: Module, description: \"mod test_mod\" }, Atom(Flag(\"test\")))",
+                ]
+            "#]],
+        )
+    }
 }