Fix Scala 3 shouldReturn/mustReturn inference for polymorphic methods

core/src/main/scala-3/org/mockito/IdiomaticStubbing.scala

@@ -3,7 +3,6 @@ package org.mockito
 import org.mockito.WhenMacroRuntime.{ AnswerActions, AnswerPFActions, RealMethod }
 import org.mockito.stubbing.ScalaOngoingStubbing
 import scala.quoted.*
-
 import org.mockito.WhenMacro
 import org.mockito.DoSomethingMacro
 
@@ -16,9 +15,13 @@ trait IdiomaticStubbing extends IdiomaticStubbingRuntime {
   val called: Called.type = Called
 
   extension [T](inline stubbing: T) {
-    transparent inline def shouldReturn: ReturnActions[T] = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
-    transparent inline def mustReturn: ReturnActions[T]   = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
-    transparent inline def returns: ReturnActions[T]      = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
+    transparent inline def shouldReturn: ReturnActions[T]                        = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
+    transparent inline def mustReturn: ReturnActions[T]                          = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
+    transparent inline def returns: ReturnActions[T]                             = WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]]
+    transparent inline infix def shouldReturn(value: T): ScalaOngoingStubbing[T] =
+      WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]](value)
+    transparent inline infix def mustReturn(value: T): ScalaOngoingStubbing[T] =
+      WhenMacro.shouldReturn[T](stubbing).asInstanceOf[ReturnActions[T]](value)
 
     transparent inline def shouldCall(crm: RealMethod.type): ScalaOngoingStubbing[T] = WhenMacro.shouldCallRealMethod[T](stubbing)(using org.scalactic.Prettifier.default)
     transparent inline def mustCall(crm: RealMethod.type): ScalaOngoingStubbing[T]   = WhenMacro.shouldCallRealMethod[T](stubbing)(using org.scalactic.Prettifier.default)

scalatest/src/test/scala/user/org/mockito/IdiomaticStubbingTest.scala

@@ -9,6 +9,12 @@ import org.scalatest.wordspec.AnyWordSpec
 import user.org.mockito.matchers.{ ValueCaseClassInt, ValueCaseClassString, ValueClass }
 import scala.collection.parallel.immutable
 
+trait PolymorphicCodec[A]
+case class PolymorphicResponse[A](body: A)
+trait PolymorphicClient {
+  def request[A](path: String)(implicit codec: PolymorphicCodec[A]): Either[String, PolymorphicResponse[A]]
+}
+
 class IdiomaticStubbingTest extends AnyWordSpec with Matchers with ArgumentMatchersSugar with IdiomaticMockitoTestSetup with IdiomaticStubbing {
 
   forAll(scenarios) { (testDouble, orgDouble, foo) =>
@@ -359,6 +365,40 @@ class IdiomaticStubbingTest extends AnyWordSpec with Matchers with ArgumentMatch
   }
 
   "mock" should {
+    "infer the type parameter for shouldReturn on a polymorphic method from the returned value" in {
+      implicit object StringCodec extends PolymorphicCodec[String]
+
+      def stubResponse[A](
+          client: PolymorphicClient,
+          response: Either[String, PolymorphicResponse[A]]
+      )(implicit codec: PolymorphicCodec[A]): org.mockito.stubbing.ScalaOngoingStubbing[Either[String, PolymorphicResponse[A]]] =
+        client.request("path")(*) shouldReturn response
+
+      val client   = mock[PolymorphicClient]
+      val expected = Right(PolymorphicResponse("ok")): Either[String, PolymorphicResponse[String]]
+
+      stubResponse(client, expected)
+
+      client.request[String]("path") shouldBe expected
+    }
+
+    "infer the type parameter for mustReturn on a polymorphic method from the returned value" in {
+      implicit object StringCodec extends PolymorphicCodec[String]
+
+      def stubResponse[A](
+          client: PolymorphicClient,
+          response: Either[String, PolymorphicResponse[A]]
+      )(implicit codec: PolymorphicCodec[A]): org.mockito.stubbing.ScalaOngoingStubbing[Either[String, PolymorphicResponse[A]]] =
+        client.request("path")(*) mustReturn response
+
+      val client   = mock[PolymorphicClient]
+      val expected = Right(PolymorphicResponse("ok")): Either[String, PolymorphicResponse[String]]
+
+      stubResponse(client, expected)
+
+      client.request[String]("path") shouldBe expected
+    }
+
     "stub a map" in {
       val mocked = mock[Map[String, String]]
       mocked(*) returns "123"