diff --git a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala index 73f7ef1c86..394dac92f9 100644 --- a/zio-http/src/test/scala/zio/http/WebSocketSpec.scala +++ b/zio-http/src/test/scala/zio/http/WebSocketSpec.scala @@ -26,6 +26,35 @@ import zio.http.ChannelEvent.{Read, Unregistered, UserEvent, UserEventTriggered} import zio.http.internal.{DynamicServer, HttpRunnableSpec, severTestLayer} object WebSocketSpec extends HttpRunnableSpec { + implicit class scopeDisconnect(scope: zio.Scope) { + def disconnect(label: String): Scope = + new Scope { + def addFinalizerExit(finalizer: Exit[Any, Any] => UIO[Any])(implicit trace: zio.Trace): UIO[Unit] = + scope.addFinalizerExit { (exit: Exit[Any, Any]) => + val warn = + ZIO + .logWarning( + s"A finalizer for layer ${label} has taken more than 1 minute to complete. Skipping this finalizer and moving onto the next one.", + ) + .delay(1.minute) + .unit + + finalizer(exit).disconnect.race(warn) + } + + def forkWith(executionStrategy: => zio.ExecutionStrategy)(implicit trace: zio.Trace): UIO[Scope.Closeable] = + scope.forkWith(executionStrategy) + } + } + implicit class layerDisconnect[I, E, O](layer: ZLayer[I, E, O]) { + def disconnect(label: String): ZLayer[I, E, O] = + ZLayer.scopedEnvironment[I] { + for { + scope <- ZIO.scope + zenv <- layer.build(scope.disconnect(label)) + } yield zenv + } + } private val websocketSpec = suite("WebsocketSpec")( test("channel events between client and server") { @@ -212,7 +241,12 @@ object WebSocketSpec extends HttpRunnableSpec { serve.as(List(websocketSpec)) } } - .provideShared(DynamicServer.live, severTestLayer, Client.default, Scope.default) @@ + .provideShared( + DynamicServer.live.disconnect("DynamicServer.live"), + severTestLayer.disconnect("serverTestLayer"), + Client.default.disconnect("Client.default"), + Scope.default, + ) @@ timeout(30 seconds) @@ diagnose(30.seconds) @@ withLiveClock @@ sequential final class MessageCollector[A](ref: Ref[List[A]], promise: Promise[Nothing, Unit]) {