3 files changed, 38 insertions, 26 deletions

diff --git a/crates/ide/src/diagnostics.rs b/crates/ide/src/diagnostics.rs
index 1c7f02763..3df73ed4f 100644
--- a/crates/ide/src/diagnostics.rs
+++ b/crates/ide/src/diagnostics.rs
@@ -10,7 +10,7 @@ mod field_shorthand;
 use std::cell::RefCell;
 
 use hir::{
-    diagnostics::{Diagnostic as _, DiagnosticSinkBuilder},
+    diagnostics::{Diagnostic as _, DiagnosticCode, DiagnosticSinkBuilder},
     Semantics,
 };
 use ide_db::base_db::SourceDatabase;
@@ -35,15 +35,23 @@ pub struct Diagnostic {
     pub severity: Severity,
     pub fix: Option<Fix>,
     pub unused: bool,
+    pub code: Option<DiagnosticCode>,
 }
 
 impl Diagnostic {
     fn error(range: TextRange, message: String) -> Self {
-        Self { message, range, severity: Severity::Error, fix: None, unused: false }
+        Self { message, range, severity: Severity::Error, fix: None, unused: false, code: None }
     }
 
     fn hint(range: TextRange, message: String) -> Self {
-        Self { message, range, severity: Severity::WeakWarning, fix: None, unused: false }
+        Self {
+            message,
+            range,
+            severity: Severity::WeakWarning,
+            fix: None,
+            unused: false,
+            code: None,
+        }
     }
 
     fn with_fix(self, fix: Option<Fix>) -> Self {
@@ -53,6 +61,10 @@ impl Diagnostic {
     fn with_unused(self, unused: bool) -> Self {
         Self { unused, ..self }
     }
+
+    fn with_code(self, code: Option<DiagnosticCode>) -> Self {
+        Self { code, ..self }
+    }
 }
 
 #[derive(Debug)]
@@ -126,7 +138,8 @@ pub(crate) fn diagnostics(
             // Override severity and mark as unused.
             res.borrow_mut().push(
                 Diagnostic::hint(sema.diagnostics_display_range(d).range, d.message())
-                    .with_unused(true),
+                    .with_unused(true)
+                    .with_code(Some(d.code())),
             );
         })
         // Only collect experimental diagnostics when they're enabled.
@@ -137,8 +150,10 @@ pub(crate) fn diagnostics(
     let mut sink = sink_builder
         // Diagnostics not handled above get no fix and default treatment.
         .build(|d| {
-            res.borrow_mut()
-                .push(Diagnostic::error(sema.diagnostics_display_range(d).range, d.message()));
+            res.borrow_mut().push(
+                Diagnostic::error(sema.diagnostics_display_range(d).range, d.message())
+                    .with_code(Some(d.code())),
+            );
         });
 
     if let Some(m) = sema.to_module_def(file_id) {
@@ -149,11 +164,15 @@ pub(crate) fn diagnostics(
 }
 
 fn diagnostic_with_fix<D: DiagnosticWithFix>(d: &D, sema: &Semantics<RootDatabase>) -> Diagnostic {
-    Diagnostic::error(sema.diagnostics_display_range(d).range, d.message()).with_fix(d.fix(&sema))
+    Diagnostic::error(sema.diagnostics_display_range(d).range, d.message())
+        .with_fix(d.fix(&sema))
+        .with_code(Some(d.code()))
 }
 
 fn warning_with_fix<D: DiagnosticWithFix>(d: &D, sema: &Semantics<RootDatabase>) -> Diagnostic {
-    Diagnostic::hint(sema.diagnostics_display_range(d).range, d.message()).with_fix(d.fix(&sema))
+    Diagnostic::hint(sema.diagnostics_display_range(d).range, d.message())
+        .with_fix(d.fix(&sema))
+        .with_code(Some(d.code()))
 }
 
 fn check_unnecessary_braces_in_use_statement(
@@ -589,6 +608,11 @@ fn test_fn() {
                             },
                         ),
                         unused: false,
+                        code: Some(
+                            DiagnosticCode(
+                                "unresolved-module",
+                            ),
+                        ),
                     },
                 ]
             "#]],
diff --git a/crates/ide/src/fn_references.rs b/crates/ide/src/fn_references.rs
index 459f201ed..5cbbe306e 100644
--- a/crates/ide/src/fn_references.rs
+++ b/crates/ide/src/fn_references.rs
@@ -1,11 +1,12 @@
 //! This module implements a methods and free functions search in the specified file.
 //! We have to skip tests, so cannot reuse file_structure module.
 
+use assists::utils::test_related_attribute;
 use hir::Semantics;
 use ide_db::RootDatabase;
 use syntax::{ast, ast::NameOwner, AstNode, SyntaxNode};
 
-use crate::{runnables::has_test_related_attribute, FileId, FileRange};
+use crate::{FileId, FileRange};
 
 pub(crate) fn find_all_methods(db: &RootDatabase, file_id: FileId) -> Vec<FileRange> {
     let sema = Semantics::new(db);
@@ -15,7 +16,7 @@ pub(crate) fn find_all_methods(db: &RootDatabase, file_id: FileId) -> Vec<FileRa
 
 fn method_range(item: SyntaxNode, file_id: FileId) -> Option<FileRange> {
     ast::Fn::cast(item).and_then(|fn_def| {
-        if has_test_related_attribute(&fn_def) {
+        if test_related_attribute(&fn_def).is_some() {
             None
         } else {
             fn_def.name().map(|name| FileRange { file_id, range: name.syntax().text_range() })
diff --git a/crates/ide/src/runnables.rs b/crates/ide/src/runnables.rs
index 2bd0e86e5..e15411777 100644
--- a/crates/ide/src/runnables.rs
+++ b/crates/ide/src/runnables.rs
@@ -1,5 +1,6 @@
 use std::fmt;
 
+use assists::utils::test_related_attribute;
 use cfg::CfgExpr;
 use hir::{AsAssocItem, Attrs, HirFileId, InFile, Semantics};
 use ide_db::RootDatabase;
@@ -156,7 +157,7 @@ fn runnable_fn(
             None => TestId::Name(name_string),
         };
 
-        if has_test_related_attribute(&fn_def) {
+        if test_related_attribute(&fn_def).is_some() {
             let attr = TestAttr::from_fn(&fn_def);
             RunnableKind::Test { test_id, attr }
         } else if fn_def.has_atom_attr("bench") {
@@ -235,20 +236,6 @@ impl TestAttr {
     }
 }
 
-/// This is a method with a heuristics to support test methods annotated with custom test annotations, such as
-/// `#[test_case(...)]`, `#[tokio::test]` and similar.
-/// Also a regular `#[test]` annotation is supported.
-///
-/// 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(crate) fn has_test_related_attribute(fn_def: &ast::Fn) -> bool {
-    fn_def
-        .attrs()
-        .filter_map(|attr| attr.path())
-        .map(|path| path.syntax().to_string().to_lowercase())
-        .any(|attribute_text| attribute_text.contains("test"))
-}
-
 const RUSTDOC_FENCE: &str = "```";
 const RUSTDOC_CODE_BLOCK_ATTRIBUTES_RUNNABLE: &[&str] =
     &["", "rust", "should_panic", "edition2015", "edition2018"];
@@ -307,7 +294,7 @@ fn has_test_function_or_multiple_test_submodules(module: &ast::Module) -> bool {
         for item in item_list.items() {
             match item {
                 ast::Item::Fn(f) => {
-                    if has_test_related_attribute(&f) {
+                    if test_related_attribute(&f).is_some() {
                         return true;
                     }
                 }