4 files changed, 35 insertions, 40 deletions

diff --git a/macros/src/expectation.rs b/macros/src/expectation.rs
index ff3d192..4fc1451 100644
--- a/macros/src/expectation.rs
+++ b/macros/src/expectation.rs
@@ -13,6 +13,7 @@ use syn::{
     GenericArgument,
     GenericParam,
     Generics,
+    ImplItemMethod,
     ItemStruct,
     Lifetime,
     Path,
@@ -20,7 +21,6 @@ use syn::{
     Receiver,
     ReturnType,
     Token,
-    TraitItemMethod,
     Type,
     TypeBareFn,
     TypePath,
@@ -60,7 +60,7 @@ pub struct Expectation
 
 impl Expectation
 {
-    pub fn new(mock: &Ident, item_method: &TraitItemMethod) -> Self
+    pub fn new(mock: &Ident, item_method: &ImplItemMethod) -> Self
     {
         let ident = create_expectation_ident(mock, &item_method.sig.ident);
 
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index ce91f87..36c6ad7 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -2,7 +2,7 @@
 use proc_macro::TokenStream;
 use proc_macro_error::{proc_macro_error, ResultExt};
 use quote::{format_ident, quote};
-use syn::{parse, TraitItem};
+use syn::{parse, ImplItem};
 
 use crate::expectation::Expectation;
 use crate::mock::Mock;
@@ -25,10 +25,11 @@ pub fn mock(input_stream: TokenStream) -> TokenStream
     let mock_mod_ident = format_ident!("__{mock_ident}");
 
     let method_items = input
+        .item_impl
         .items
         .into_iter()
         .filter_map(|item| match item {
-            TraitItem::Method(item_method) => Some(item_method),
+            ImplItem::Method(item_method) => Some(item_method),
             _ => None,
         })
         .collect::<Vec<_>>();
diff --git a/macros/src/mock.rs b/macros/src/mock.rs
index 88d3434..1b6dd67 100644
--- a/macros/src/mock.rs
+++ b/macros/src/mock.rs
@@ -16,7 +16,6 @@ use syn::{
     ReturnType,
     Signature,
     Token,
-    TraitItemMethod,
     Type,
     TypePath,
     TypeReference,
@@ -43,18 +42,15 @@ use crate::util::create_unit_type_tuple;
 pub struct Mock
 {
     ident: Ident,
-    mocked_trait: TypePath,
+    mocked_trait: Path,
     expectations_fields: Vec<ExpectationsField>,
-    item_methods: Vec<TraitItemMethod>,
+    item_methods: Vec<ImplItemMethod>,
 }
 
 impl Mock
 {
-    pub fn new(
-        ident: Ident,
-        mocked_trait: TypePath,
-        item_methods: &[TraitItemMethod],
-    ) -> Self
+    pub fn new(ident: Ident, mocked_trait: Path, item_methods: &[ImplItemMethod])
+        -> Self
     {
         let expectations_fields = item_methods
             .iter()
@@ -179,7 +175,7 @@ impl ToTokens for ExpectationsField
     }
 }
 
-fn create_mock_function(item_method: TraitItemMethod) -> ImplItemMethod
+fn create_mock_function(item_method: ImplItemMethod) -> ImplItemMethod
 {
     let func_ident = &item_method.sig.ident;
 
@@ -247,7 +243,7 @@ fn create_mock_function(item_method: TraitItemMethod) -> ImplItemMethod
     }
 }
 
-fn create_expect_function(mock: &Ident, item_method: &TraitItemMethod) -> ImplItemMethod
+fn create_expect_function(mock: &Ident, item_method: &ImplItemMethod) -> ImplItemMethod
 {
     let signature = Signature::new(
         format_ident!("expect_{}", item_method.sig.ident),
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,
         })
     }
 }