@@ -302,7 +302,8 @@ fn test_custom_fn_attribute() {
302302 use std:: env;
303303 use std:: fs;
304304 use std:: path:: Path ;
305- use syn:: { parse_file, File , Item , ItemFn } ;
305+ use syn:: visit:: Visit ;
306+ use syn:: { parse_file, File , ForeignItem , Item , ItemForeignMod } ;
306307
307308 let out_dir =
308309 std:: env:: var ( "OUT_DIR" ) . expect ( "OUT_DIR environment variable not set" ) ;
@@ -312,27 +313,42 @@ fn test_custom_fn_attribute() {
312313 let syntax_tree: File =
313314 parse_file ( & file_content) . expect ( "Failed to parse test.rs" ) ;
314315
315- let mut found_coord = false ;
316- let mut has_must_use = false ;
317-
318- for item in syntax_tree. items {
319- if let Item :: Fn ( item_fn) = item {
320- if item_fn. sig . ident == "coord" {
321- found_coord = true ;
322- has_must_use = item_fn
323- . attrs
324- . iter ( )
325- . any ( |attr| attr. path ( ) . is_ident ( "must_use" ) ) ;
316+ struct FunctionVisitor {
317+ found_coord : bool ,
318+ has_must_use : bool ,
319+ }
320+
321+ impl < ' ast > Visit < ' ast > for FunctionVisitor {
322+ fn visit_item_foreign_mod (
323+ & mut self ,
324+ foreign_mod : & ' ast ItemForeignMod ,
325+ ) {
326+ for foreign_item in & foreign_mod. items {
327+ if let ForeignItem :: Fn ( item_fn) = foreign_item {
328+ if item_fn. sig . ident == "coord" {
329+ self . found_coord = true ;
330+ self . has_must_use = item_fn
331+ . attrs
332+ . iter ( )
333+ . any ( |attr| attr. path ( ) . is_ident ( "must_use" ) ) ;
334+ }
335+ }
326336 }
327337 }
328338 }
329339
340+ let mut visitor = FunctionVisitor {
341+ found_coord : false ,
342+ has_must_use : false ,
343+ } ;
344+ visitor. visit_file ( & syntax_tree) ;
345+
330346 assert ! (
331- found_coord,
347+ visitor . found_coord,
332348 "The function 'coord' was not found in the source."
333349 ) ;
334350 assert ! (
335- has_must_use,
351+ visitor . has_must_use,
336352 "The function 'coord' does not have the #[must_use] attribute."
337353 ) ;
338354}
0 commit comments