add new spec

build.sbt

@@ -62,3 +62,9 @@ lazy val root = project
       "org.scala-lang" %% "scala3-compiler" % scala3Version % "provided"
     )
   )
+
+/////////////////////
+// Command Aliases //
+/////////////////////
+
+addCommandAlias("prepare", "scalafmtAll; scalafixAll")

src/main/scala/stubby/package.scala

@@ -7,9 +7,21 @@ import scala.jdk.CollectionConverters.*
 case class MethodId(name: String)
 
 class ServicePartiallyApplied[Service]:
-  inline def apply[Output](inline select: Service => Output)(inline result: Output): URIO[Stubbed[Service], Unit] =
+  inline def apply[Output](inline select: Service => Output)(
+      inline result: Output
+  ): URIO[Stubbed[Service], Unit] =
     ${ Macros.stubImpl('select, 'result) }
 
+  // inline def apply[Output, A](inline select: Service => (A) => Output)(
+  //     inline handler: A => Output
+  // ): URIO[Stubbed[Service], Unit] =
+  //   ${ Macros.stubImpl('select, 'handler) }
+
+  // inline def apply[Output, A, B](inline select: Service => (A, B) => Output)(
+  //     inline handler: (A, B) => Output
+  // ): URIO[Stubbed[Service], Unit] =
+  //   ${ Macros.stubImpl('select, 'handler) }
+
 def any[A]: A = ???
 
 inline def stub[Service]                                        = ServicePartiallyApplied[Service]()
@@ -32,24 +44,57 @@ trait Stubbed[A]:
     response.asInstanceOf[A]
 
 object Stubbed:
-  def insert[Service: Tag](
+  def insert[Service: Tag, Value](
       methodId: MethodId,
-      response: Any
+      response: Value
   ): URIO[Stubbed[Service], Unit] =
     ZIO.serviceWithZIO(_.insert(methodId, response))
 
 object Macros:
 
+  def getMethod[Service: Type, Output: Type](using Quotes)(select: Expr[Service => Output]): quotes.reflect.Symbol =
+    import quotes.reflect.*
+    select.asTerm.underlyingArgument match
+
+      // def methodName(arg: Int): Result = ???
+      // stub[Service](_.methodName(any))
+      case Lambda(_, Apply(select @ Select(_, methodName), _)) =>
+        select.symbol
+
+      // def methodName: Result = ???
+      // stub[Service](_.methodName)
+      case Lambda(_, select @ Select(_, methodName)) =>
+        select.symbol
+
+      // def methodName(arg: Int): Result = ???
+      // stub[Service](_.methodName)
+      case Lambda(args, Lambda(_, Apply(select @ Select(_, methodName), _))) =>
+        select.symbol
+
+      case _ =>
+        // TODO: better error message
+        report.errorAndAbort(s"Invalid selector: ${select.show}")
+
   def stubImpl[Service: Type, Output: Type](
       select: Expr[Service => Output],
       result: Expr[Output]
   )(using Quotes): Expr[URIO[Stubbed[Service], Unit]] =
     import quotes.reflect.*
 
-    select.asTerm.underlyingArgument match
-      case Lambda(args, body @ Apply(Select(_, methodName), _)) =>
+    val method = getMethod(select)
+    method.termRef.widenTermRefByName.returnType.asType match
+      case '[t] =>
+        val resultExpr =
+          try result.asExprOf[t]
+          catch
+            case _: Exception =>
+              report.errorAndAbort(
+                s"Expected ${Type.show[t]} but got ${result.asTerm.tpe.widen.show}",
+                result.asTerm.pos
+              )
+
         '{
-          Stubbed.insert[Service](MethodId(${ Expr(methodName) }), $result)
+          Stubbed.insert[Service, t](MethodId(${ Expr(method.name) }), $resultExpr)
         }
 
   def stubbedImpl[Service: Type](using Quotes): Expr[ULayer[Service & Stubbed[Service]]] =

src/test/scala/StubbySpec.scala

@@ -0,0 +1,123 @@
+import zio.test.*
+import zio.*
+
+import stubby.*
+import scala.compiletime.testing.typeCheckErrors
+
+case class User(id: Int, name: String)
+
+trait ExampleService:
+  def effectWithParams(id: Int): Task[User]
+  def effectWithoutParams: Task[Int]
+  def pureWithParams(id: Int): Int
+  def pureWithoutParams: Int
+
+  def overloadedEffect(id: Int): Task[User]
+  def overloadedEffect(name: String): Task[User]
+
+object ExampleService:
+  def effectWithParams(id: Int): ZIO[ExampleService, Throwable, User] =
+    ZIO.serviceWithZIO[ExampleService](_.effectWithParams(id))
+  def effectWithoutParams: ZIO[ExampleService, Throwable, Int] =
+    ZIO.serviceWithZIO[ExampleService](_.effectWithoutParams)
+  def pureWithParams(id: Int): ZIO[ExampleService, Nothing, Int] =
+    ZIO.serviceWith[ExampleService](_.pureWithParams(id))
+  def pureWithoutParams: ZIO[ExampleService, Nothing, Int] =
+    ZIO.serviceWith[ExampleService](_.pureWithoutParams)
+  def overloadedEffect(id: Int): ZIO[ExampleService, Throwable, User] =
+    ZIO.serviceWithZIO[ExampleService](_.overloadedEffect(id))
+  def overloadedEffect(name: String): ZIO[ExampleService, Throwable, User] =
+    ZIO.serviceWithZIO[ExampleService](_.overloadedEffect(name))
+
+object StubbySpec extends ZIOSpecDefault:
+  val spec =
+    suiteAll("StubbySpec") {
+
+      suiteAll("stub") {
+
+        suiteAll("effectful method with params") {
+
+          test("succeeds when leaving off the params") {
+            for
+              _ <- stub[ExampleService](_.effectWithParams) {
+                     ZIO.succeed(User(1, "John Doe"))
+                   }
+              result <- ExampleService.effectWithParams(1)
+            yield assertTrue(result == User(1, "John Doe"))
+          }
+
+          test("succeeds when providing any params") {
+            for
+              _ <- stub[ExampleService](_.effectWithParams(any)) {
+                     ZIO.succeed(User(1, "John Doe"))
+                   }
+              result <- ExampleService.effectWithParams(1)
+            yield assertTrue(result == User(1, "John Doe"))
+          }
+
+          // test("succeeds when providing a handler") {
+          //   for
+          //     _ <- stub[ExampleService](_.effectWithParams) { arg =>
+          //            ZIO.succeed(User(arg, "John Doe"))
+          //          }
+          //     result <- ExampleService.effectWithParams(1)
+          //   yield assertTrue(result == User(1, "John Doe"))
+          // }
+        }
+
+        suiteAll("effectful method without params") {
+
+          test("succeeds when calling") {
+            for
+              _ <- stub[ExampleService](_.effectWithoutParams) {
+                     ZIO.succeed(12)
+                   }
+              result <- ExampleService.effectWithoutParams
+            yield assertTrue(result == 12)
+          }
+
+        }
+
+        suiteAll("overloaded method") {
+
+          test("succeeds when calling with id") {
+            for
+              _ <- stub[ExampleService](_.overloadedEffect(any[Int])) {
+                     ZIO.succeed(User(1, "John Doe"))
+                   }
+              result <- ExampleService.overloadedEffect(1)
+            yield assertTrue(result == User(1, "John Doe"))
+          }
+
+          test("succeeds when calling with name") {
+            for
+              _ <- stub[ExampleService](_.overloadedEffect(any[String])) {
+                     ZIO.succeed(User(1, "John Doe"))
+                   }
+              result <- ExampleService.overloadedEffect("John Doe")
+            yield assertTrue(result == User(1, "John Doe"))
+          }
+
+        }
+
+      }
+
+      suiteAll("error messages") {
+
+        test("providing the wrong type") {
+          val errors = typeCheckErrors {
+            """
+            stub[ExampleService](_.effectWithParams) {
+              12
+            }
+            """
+          }
+
+          assertTrue(errors.head.message == "Expected zio.Task[User] but got scala.Int")
+        }
+
+      }
+
+    }.provide(
+      stubbed[ExampleService]
+    )