3 files changed, 21 insertions, 31 deletions
diff --git a/src/libsyntax_ext/lib.rs b/src/libsyntax_ext/lib.rs
index 62530f4fe7b..da0f8ca6da0 100644
--- a/src/libsyntax_ext/lib.rs
+++ b/src/libsyntax_ext/lib.rs
@@ -72,6 +72,7 @@ pub fn register_builtins(resolver: &mut dyn syntax::ext::base::Resolver,
     let mut register = |name, ext| {
         resolver.add_builtin(ast::Ident::with_empty_ctxt(name), Lrc::new(ext));
     };
+
     macro_rules! register {
         ($( $name:ident: $f:expr, )*) => { $(
             register(sym::$name, SyntaxExtension::default(
@@ -125,24 +126,31 @@ pub fn register_builtins(resolver: &mut dyn syntax::ext::base::Resolver,
         trace_macros: trace_macros::expand_trace_macros,
     }
 
+    let allow_internal_unstable = Some([sym::test, sym::rustc_attrs][..].into());
     register(sym::test_case, SyntaxExtension {
         stability: Some(Stability::unstable(
             sym::custom_test_frameworks,
             Some(Symbol::intern(EXPLAIN_CUSTOM_TEST_FRAMEWORKS)),
             50297,
         )),
+        allow_internal_unstable: allow_internal_unstable.clone(),
         ..SyntaxExtension::default(
             SyntaxExtensionKind::LegacyAttr(Box::new(test_case::expand)), edition
         )
     });
-    register(sym::test, SyntaxExtension::default(
-        SyntaxExtensionKind::LegacyAttr(Box::new(test::expand_test)), edition
-    ));
-    register(sym::bench, SyntaxExtension::default(
-        SyntaxExtensionKind::LegacyAttr(Box::new(test::expand_bench)), edition
-    ));
-
-    // format_args uses `unstable` things internally.
+    register(sym::test, SyntaxExtension {
+        allow_internal_unstable: allow_internal_unstable.clone(),
+        ..SyntaxExtension::default(
+            SyntaxExtensionKind::LegacyAttr(Box::new(test::expand_test)), edition
+        )
+    });
+    register(sym::bench, SyntaxExtension {
+        allow_internal_unstable,
+        ..SyntaxExtension::default(
+            SyntaxExtensionKind::LegacyAttr(Box::new(test::expand_bench)), edition
+        )
+    });
+
     let allow_internal_unstable = Some([sym::fmt_internals][..].into());
     register(sym::format_args, SyntaxExtension {
         allow_internal_unstable: allow_internal_unstable.clone(),
diff --git a/src/libsyntax_ext/test.rs b/src/libsyntax_ext/test.rs
index 7f4f1fd7e01..c5c5ef57b31 100644
--- a/src/libsyntax_ext/test.rs
+++ b/src/libsyntax_ext/test.rs
@@ -3,13 +3,12 @@
 
 use syntax::ext::base::*;
 use syntax::ext::build::AstBuilder;
-use syntax::ext::hygiene::{Mark, SyntaxContext};
+use syntax::ext::hygiene::SyntaxContext;
 use syntax::attr;
 use syntax::ast;
 use syntax::print::pprust;
 use syntax::symbol::{Symbol, sym};
 use syntax_pos::Span;
-use syntax::source_map::{ExpnInfo, ExpnKind};
 use std::iter;
 
 pub fn expand_test(
@@ -60,15 +59,8 @@ pub fn expand_test_or_bench(
         return vec![Annotatable::Item(item)];
     }
 
-    let (sp, attr_sp) = {
-        let mark = Mark::fresh(Mark::root());
-        mark.set_expn_info(ExpnInfo::with_unstable(
-            ExpnKind::MacroAttribute(sym::test), attr_sp, cx.parse_sess.edition,
-            &[sym::rustc_attrs, sym::test],
-        ));
-        (item.span.with_ctxt(SyntaxContext::empty().apply_mark(mark)),
-         attr_sp.with_ctxt(SyntaxContext::empty().apply_mark(mark)))
-    };
+    let ctxt = SyntaxContext::empty().apply_mark(cx.current_expansion.mark);
+    let (sp, attr_sp) = (item.span.with_ctxt(ctxt), attr_sp.with_ctxt(ctxt));
 
     // Gensym "test" so we can extern crate without conflicting with any local names
     let test_id = cx.ident_of("test").gensym();
diff --git a/src/libsyntax_ext/test_case.rs b/src/libsyntax_ext/test_case.rs
index 5f88e2e098a..af2cf42e04b 100644
--- a/src/libsyntax_ext/test_case.rs
+++ b/src/libsyntax_ext/test_case.rs
@@ -11,12 +11,11 @@
 
 use syntax::ext::base::*;
 use syntax::ext::build::AstBuilder;
-use syntax::ext::hygiene::{Mark, SyntaxContext};
+use syntax::ext::hygiene::SyntaxContext;
 use syntax::ast;
 use syntax::source_map::respan;
 use syntax::symbol::sym;
 use syntax_pos::Span;
-use syntax::source_map::{ExpnInfo, ExpnKind};
 
 pub fn expand(
     ecx: &mut ExtCtxt<'_>,
@@ -26,17 +25,8 @@ pub fn expand(
 ) -> Vec<Annotatable> {
     if !ecx.ecfg.should_test { return vec![]; }
 
-    let sp = {
-        let mark = Mark::fresh(Mark::root());
-        mark.set_expn_info(ExpnInfo::with_unstable(
-            ExpnKind::MacroAttribute(sym::test_case), attr_sp, ecx.parse_sess.edition,
-            &[sym::test, sym::rustc_attrs],
-        ));
-        attr_sp.with_ctxt(SyntaxContext::empty().apply_mark(mark))
-    };
-
+    let sp = attr_sp.with_ctxt(SyntaxContext::empty().apply_mark(ecx.current_expansion.mark));
     let mut item = anno_item.expect_item();
-
     item = item.map(|mut item| {
         item.vis = respan(item.vis.span, ast::VisibilityKind::Public);
         item.ident = item.ident.gensym();