@@ -59,9 +59,11 @@ pub struct TestMetadata {
5959pub enum ExpectedRevert {
6060 /// Test should revert with any data.
6161Any ,
62- // Future phases:
63- // ExactData(Vec<u8>),
64- // Selector([u8; 4]),
62+ /// Test should revert with data starting with the given 4-byte selector.
63+ Selector ( [ u8 ; 4 ] ) ,
64+ /// Test should revert with a Panic(uint256) whose code matches.
65+ /// Stores the expected full revert payload: selector (4 bytes) + ABI-encoded code (32 bytes).
66+ PanicCode ( Vec < u8 > ) ,
6567}
6668
6769/// Output returned by `emit_test_module_yul`.
@@ -860,19 +862,15 @@ fn collect_test_infos(
860862 } ;
861863 let attrs = ItemKind :: from ( hir_func) . attrs ( db) ?;
862864 let test_attr = attrs. get_attr ( db, "test" ) ?;
863-
864- // Check for #[test(should_revert)]
865- let expected_revert = if test_attr. has_arg ( db, "should_revert" ) {
866- Some ( ExpectedRevert :: Any )
867- } else {
868- None
869- } ;
870-
871865 let hir_name = hir_func
872866 . name ( db)
873867 . to_opt ( )
874868 . map ( |n| n. data ( db) . to_string ( ) )
875869 . unwrap_or_else ( || "<anonymous>" . to_string ( ) ) ;
870+ let expected_revert = match parse_expected_revert ( db, & hir_name, test_attr) {
871+ Ok ( expected_revert) => expected_revert,
872+ Err ( err) => return Some ( Err ( EmitModuleError :: Yul ( YulError :: Unsupported ( err) ) ) ) ,
873+ } ;
876874 // Check for #[test(balance = N)]
877875 let initial_balance = match parse_test_balance_arg ( db, & hir_name, test_attr) {
878876 Ok ( balance) => balance,
@@ -931,6 +929,121 @@ fn parse_test_balance_arg<'db>(
931929
932930 Ok ( None )
933931}
932+ /// Parses the expected revert behavior from a `#[test(...)]` attribute.
933+ ///
934+ /// Supported forms:
935+ /// - `#[test(should_revert)]` — any revert
936+ /// - `#[test(should_revert, selector = 0x4e487b71)]` — revert with matching 4-byte selector
937+ /// - `#[test(should_revert, panic = 0x11)]` — revert with Panic(uint256) and matching code
938+ pub fn parse_expected_revert < ' db > (
939+ db : & ' db dyn HirDb ,
940+ test_name : & str ,
941+ test_attr : & hir:: hir_def:: attr:: NormalAttr < ' db > ,
942+ ) -> Result < Option < ExpectedRevert > , String > {
943+ let should_revert = test_attr. has_arg ( db, "should_revert" ) ;
944+ let has_panic = has_test_attr_key ( db, test_attr, "panic" ) ;
945+ let has_selector = has_test_attr_key ( db, test_attr, "selector" ) ;
946+
947+ if !should_revert {
948+ if has_panic && has_selector {
949+ return Err ( format ! (
950+ "invalid #[test] function `{test_name}`: `panic = ...` and `selector = ...` require `should_revert`"
951+ ) ) ;
952+ }
953+ if has_panic {
954+ return Err ( format ! (
955+ "invalid #[test] function `{test_name}`: `panic = ...` requires `should_revert`"
956+ ) ) ;
957+ }
958+ if has_selector {
959+ return Err ( format ! (
960+ "invalid #[test] function `{test_name}`: `selector = ...` requires `should_revert`"
961+ ) ) ;
962+ }
963+ return Ok ( None ) ;
964+ }
965+
966+ let panic = parse_test_attr_int_arg ( db, test_name, test_attr, "panic" , "u256" , 32 ) ?;
967+ let selector = parse_test_attr_int_arg ( db, test_name, test_attr, "selector" , "u32" , 4 ) ?;
968+
969+ if panic. is_some ( ) && selector. is_some ( ) {
970+ return Err ( format ! (
971+ "invalid #[test] function `{test_name}`: #[test(should_revert)] cannot combine `panic = ...` and `selector = ...`"
972+ ) ) ;
973+ }
974+
975+ if let Some ( code) = panic {
976+ let panic_selector: [ u8 ; 4 ] = [ 0x4e , 0x48 , 0x7b , 0x71 ] ;
977+ let mut payload = Vec :: with_capacity ( 36 ) ;
978+ payload. extend_from_slice ( & panic_selector) ;
979+ // ABI-encode the uint256 code: pad to 32 bytes big-endian
980+ let code_bytes = code. to_bytes_be ( ) ;
981+ let mut padded = [ 0u8 ; 32 ] ;
982+ let start = 32 - code_bytes. len ( ) ;
983+ padded[ start..] . copy_from_slice ( & code_bytes) ;
984+ payload. extend_from_slice ( & padded) ;
985+ return Ok ( Some ( ExpectedRevert :: PanicCode ( payload) ) ) ;
986+ }
987+
988+ if let Some ( sel) = selector {
989+ let bytes = sel. to_bytes_be ( ) ;
990+ let mut selector = [ 0u8 ; 4 ] ;
991+ let start = 4 - bytes. len ( ) ;
992+ selector[ start..] . copy_from_slice ( & bytes) ;
993+ return Ok ( Some ( ExpectedRevert :: Selector ( selector) ) ) ;
994+ }
995+
996+ Ok ( Some ( ExpectedRevert :: Any ) )
997+ }
998+
999+ fn has_test_attr_key < ' db > (
1000+ db : & ' db dyn HirDb ,
1001+ test_attr : & hir:: hir_def:: attr:: NormalAttr < ' db > ,
1002+ key : & str ,
1003+ ) -> bool {
1004+ test_attr
1005+ . args
1006+ . iter ( )
1007+ . any ( |arg| arg. key_str ( db) == Some ( key) )
1008+ }
1009+
1010+ fn parse_test_attr_int_arg < ' db > (
1011+ db : & ' db dyn HirDb ,
1012+ test_name : & str ,
1013+ test_attr : & hir:: hir_def:: attr:: NormalAttr < ' db > ,
1014+ key : & str ,
1015+ type_name : & str ,
1016+ max_bytes : usize ,
1017+ ) -> Result < Option < BigUint > , String > {
1018+ for arg in & test_attr. args {
1019+ if arg. key_str ( db) != Some ( key) {
1020+ continue ;
1021+ }
1022+
1023+ let Some ( value) = arg. value . as_ref ( ) else {
1024+ return Err ( format ! (
1025+ "invalid #[test] function `{test_name}`: #[test(should_revert, {key} = ...)] expects an integer literal"
1026+ ) ) ;
1027+ } ;
1028+ let hir:: hir_def:: attr:: AttrArgValue :: Lit ( hir:: hir_def:: LitKind :: Int ( int_id) ) = value
1029+ else {
1030+ return Err ( format ! (
1031+ "invalid #[test] function `{test_name}`: #[test(should_revert, {key} = ...)] expects an integer literal"
1032+ ) ) ;
1033+ } ;
1034+
1035+ let value = int_id. data ( db) . clone ( ) ;
1036+ if value. to_bytes_be ( ) . len ( ) > max_bytes {
1037+ return Err ( format ! (
1038+ "invalid #[test] function `{test_name}`: #[test(should_revert, {key} = ...)] must fit in {type_name}"
1039+ ) ) ;
1040+ }
1041+ return Ok ( Some ( value) ) ;
1042+ }
1043+
1044+ Ok ( None )
1045+ }
1046+
9341047fn test_info_matches_filter ( test : & TestInfo , filter : Option < & str > ) -> bool {
9351048 let Some ( pattern) = filter else {
9361049 return true ;
0 commit comments