1 files changed, 24 insertions, 26 deletions
diff --git a/macros/src/mock_input.rs b/macros/src/mock_input.rs
index 379c342..590ca0c 100644
--- a/macros/src/mock_input.rs
+++ b/macros/src/mock_input.rs
@@ -1,18 +1,19 @@
 use syn::parse::{Parse, ParseStream};
-use syn::{braced, Ident, Token, TraitItem, TypePath, WhereClause};
+use syn::spanned::Spanned;
+use syn::{braced, Ident, ItemImpl, Path, Type, TypePath, WhereClause};
 
 pub struct MockInput
 {
     pub mock: Ident,
-    pub mocked_trait: TypePath,
-    pub items: Vec<TraitItem>,
+    pub mocked_trait: Path,
+    pub item_impl: ItemImpl,
 }
 
 impl Parse for MockInput
 {
     fn parse(input: ParseStream) -> Result<Self, syn::Error>
     {
-        let mock = input.parse()?;
+        let mock = input.parse::<Ident>()?;
 
         let _generics = input.parse::<Option<WhereClause>>()?;
 
@@ -20,32 +21,29 @@ impl Parse for MockInput
 
         let _brace = braced!(_braced_content in input);
 
-        input.parse::<Token![impl]>()?;
-
-        let mocked_trait = input.parse()?;
-
-        input.parse::<Token![for]>()?;
-
-        let impl_target = input.parse::<Ident>()?;
-
-        if impl_target != mock {
-            return Err(input.error("Expected this to be the mock"));
-        }
-
-        let content;
-
-        braced!(content in input);
-
-        let mut items = Vec::new();
-
-        while !content.is_empty() {
-            items.push(content.parse()?);
+        let item_impl = input.parse::<ItemImpl>()?;
+
+        let Some((_, mocked_trait, _)) = &item_impl.trait_ else {
+            return Err(syn::Error::new(item_impl.impl_token.span, "Impl must be of a trait"));
+        };
+
+        if !matches!(
+            item_impl.self_ty.as_ref(),
+            Type::Path(TypePath {
+                qself: None,
+                path
+            }) if path.is_ident(&mock)
+        ) {
+            return Err(syn::Error::new(
+                item_impl.self_ty.span(),
+                "Expected this to be the mock",
+            ));
         }
 
         Ok(Self {
             mock,
-            mocked_trait,
-            items,
+            mocked_trait: mocked_trait.clone(),
+            item_impl,
         })
     }
 }