Add context method to request context

zio · Apr 25, 2024 · b870ab7 · b870ab7
1 parent 2229dcc
commit b870ab7
Show file tree

Hide file tree

Showing 4 changed files with 382 additions and 8 deletions.
diff --git a/zio-http/jvm/src/test/scala/zio/http/HandlerAspectSpec.scala b/zio-http/jvm/src/test/scala/zio/http/HandlerAspectSpec.scala
@@ -18,7 +18,7 @@ object HandlerAspectSpec extends ZIOSpecDefault {
       test("HandlerAspect with context can eliminate environment type partially") {
         val handlerAspect = HandlerAspect.interceptIncomingHandler(handler((req: Request) => (req, req.headers.size)))
         val handler0 = handler { (_: Request) =>
-          ZIO.service[Boolean] *> ZIO.serviceWith[Int](i => Response.text(i.toString))
+          withContext((_: Boolean, i: Int) => Response.text(i.toString))
           //leftover type is only needed in Scala 2
           //can't be infix because of Scala 3
         }.@@[Boolean](handlerAspect)
@@ -29,10 +29,10 @@ object HandlerAspectSpec extends ZIOSpecDefault {
       },
       test("HandlerAspect with context can eliminate environment type partially while requiring an additional environment") {
         val handlerAspect: HandlerAspect[String, Int] = HandlerAspect.interceptIncomingHandler {
-          handler((req: Request) => ZIO.serviceWith[String](s => (req.withBody(Body.fromString(s)), req.headers.size)))
+          handler((req: Request) => withContext((s: String) => (req.withBody(Body.fromString(s)), req.headers.size)))
         }
         val handler0: Handler[Boolean with String, Response, Request, Response] = handler { (r: Request) =>
-          ZIO.service[Boolean] *> ZIO.serviceWithZIO[Int] { i =>
+          ZIO.service[Boolean] *> withContext{ (i: Int) =>
             for {
               body <- r.body.asString.orDie
             } yield Response.text(s"$i $body")

diff --git a/zio-http/jvm/src/test/scala/zio/http/internal/middlewares/AuthSpec.scala b/zio-http/jvm/src/test/scala/zio/http/internal/middlewares/AuthSpec.scala
@@ -19,7 +19,7 @@ package zio.http.internal.middlewares
 import zio.Config.Secret
 import zio.test.Assertion._
 import zio.test._
-import zio.{Ref, ZEnvironment, ZIO}
+import zio.{Ref, ZIO}
 
 import zio.http._
 import zio.http.internal.HttpAppTestExtensions
@@ -74,7 +74,7 @@ object AuthSpec extends ZIOHttpSpec with HttpAppTestExtensions {
       },
       test("Extract username via context") {
         val app = (Handler.fromFunctionZIO[Request](_ =>
-          ZIO.serviceWith[AuthContext](c => Response.text(c.value)),
+          withContext((c: AuthContext) => Response.text(c.value)),
         ) @@ basicAuthContextM).merge.mapZIO(_.body.asString)
         assertZIO(app.runZIO(Request.get(URL.empty).copy(headers = successBasicHeader)))(equalTo("user"))
       },
@@ -109,14 +109,13 @@ object AuthSpec extends ZIOHttpSpec with HttpAppTestExtensions {
         val secureRoutes = Routes(
           Method.GET / "a" -> handler((_: Request) => ZIO.serviceWith[AuthContext](ctx => Response.text(ctx.value))),
           Method.GET / "b" / int("id")      -> handler((id: Int, _: Request) =>
-            ZIO.serviceWith[AuthContext](ctx => Response.text(s"for id: $id: ${ctx.value}")),
+            withContext((ctx: AuthContext) => Response.text(s"for id: $id: ${ctx.value}")),
           ),
           Method.GET / "c" / string("name") -> handler((name: String, _: Request) =>
-            ZIO.serviceWith[AuthContext](ctx => Response.text(s"for name: $name: ${ctx.value}")),
+            withContext((ctx: AuthContext) => Response.text(s"for name: $name: ${ctx.value}")),
           ),
           // Needs version of @@ that removes the context from the environment
         ) @@ basicAuthContextM
-        // Just a prove that the aspect can require an environment. Does nothing.
         val app          = secureRoutes
         for {
           s1     <- app.runZIO(Request.get(URL(Path.root / "a")).copy(headers = successBasicHeader))