From c8e28757acf66b97777184a3415fed5637dd57be Mon Sep 17 00:00:00 2001
From: Audun Halland <audun.halland@pm.me>
Date: Wed, 27 Mar 2024 01:35:18 +0100
Subject: [PATCH] fix: unmock with async fn when the trait fn uses RPIT future

---
 CHANGELOG.md                         |  2 ++
 tests/it/unmock.rs                   | 15 +++++++++++++++
 unimock_macros/src/unimock/method.rs |  2 +-
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 462643a..f25e558 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
+### Fixed
+- Unmocking using `async fn` when the trait fn has an `-> impl Future` signature.
 
 ## [0.6.0] - 2024-03-25
 ### Changed
diff --git a/tests/it/unmock.rs b/tests/it/unmock.rs
index 54d1b71..069f903 100644
--- a/tests/it/unmock.rs
+++ b/tests/it/unmock.rs
@@ -150,6 +150,21 @@ mod unmock_async {
     }
 }
 
+mod impl_future_in_trait_unmock_with_async_fn {
+    use core::future::Future;
+
+    use super::*;
+
+    #[unimock(unmock_with=[foo])]
+    trait Async {
+        fn foo(&self) -> impl Future<Output = i32>;
+    }
+
+    async fn foo(_: &impl core::any::Any) -> i32 {
+        42
+    }
+}
+
 mod unmock_with_custom_args {
     use super::*;
 
diff --git a/unimock_macros/src/unimock/method.rs b/unimock_macros/src/unimock/method.rs
index 6bacef0..10ec22d 100644
--- a/unimock_macros/src/unimock/method.rs
+++ b/unimock_macros/src/unimock/method.rs
@@ -70,7 +70,7 @@ impl<'t> MockMethod<'t> {
         if self.method.sig.asyncness.is_some()
             || matches!(
                 self.output_structure.wrapping,
-                output::OutputWrapping::AssociatedFuture(_)
+                output::OutputWrapping::AssociatedFuture(_) | output::OutputWrapping::RpitFuture(_)
             )
         {
             Some(DotAwait)