From e61e3968d7bfd14b07f69813ec3c0e1cacc10d5d Mon Sep 17 00:00:00 2001
From: kciesielski <kciesielski@users.noreply.github.com>
Date: Wed, 22 Nov 2023 13:46:12 +0100
Subject: [PATCH] Fix tests

---
 .../armeria/zio/ArmeriaZioServerTest.scala      |  2 +-
 .../server/vertx/zio/ZioVertxServerTest.scala   | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/server/armeria-server/zio1/src/test/scala/sttp/tapir/server/armeria/zio/ArmeriaZioServerTest.scala b/server/armeria-server/zio1/src/test/scala/sttp/tapir/server/armeria/zio/ArmeriaZioServerTest.scala
index a3c9026586..7d1b356c84 100644
--- a/server/armeria-server/zio1/src/test/scala/sttp/tapir/server/armeria/zio/ArmeriaZioServerTest.scala
+++ b/server/armeria-server/zio1/src/test/scala/sttp/tapir/server/armeria/zio/ArmeriaZioServerTest.scala
@@ -19,6 +19,6 @@ class ArmeriaZioServerTest extends TestSuite {
 
     new AllServerTests(createServerTest, interpreter, backend, basic = false, options = false).tests() ++
       new ServerBasicTests(createServerTest, interpreter, supportsUrlEncodedPathSegments = false).tests() ++
-      new ServerStreamingTests(createServerTest, ZioStreams).tests()
+      new ServerStreamingTests(createServerTest, maxLengthSupported = false).tests(ZioStreams)(_ => Task.unit)
   }
 }
diff --git a/server/vertx-server/zio/src/test/scala/sttp/tapir/server/vertx/zio/ZioVertxServerTest.scala b/server/vertx-server/zio/src/test/scala/sttp/tapir/server/vertx/zio/ZioVertxServerTest.scala
index 908e762c26..ea448196fd 100644
--- a/server/vertx-server/zio/src/test/scala/sttp/tapir/server/vertx/zio/ZioVertxServerTest.scala
+++ b/server/vertx-server/zio/src/test/scala/sttp/tapir/server/vertx/zio/ZioVertxServerTest.scala
@@ -1,18 +1,19 @@
 package sttp.tapir.server.vertx.zio
 
+import _root_.zio.stream.ZStream
+import _root_.zio.{Task, ZIO}
 import cats.effect.{IO, Resource}
 import io.vertx.core.Vertx
+import org.scalatest.OptionValues
+import org.scalatest.matchers.should.Matchers._
 import sttp.capabilities.zio.ZioStreams
+import sttp.client3.basicRequest
 import sttp.monad.MonadError
+import sttp.tapir._
 import sttp.tapir.server.tests._
 import sttp.tapir.tests.{Test, TestSuite}
-import sttp.tapir._
-import _root_.zio.{Task, ZIO}
-import _root_.zio.stream.ZStream
-import org.scalatest.OptionValues
-import sttp.client3.basicRequest
 import sttp.tapir.ztapir.RIOMonadError
-import org.scalatest.matchers.should.Matchers._
+import zio.stream.ZSink
 
 class ZioVertxServerTest extends TestSuite with OptionValues {
   def vertxResource: Resource[IO, Vertx] =
@@ -32,6 +33,8 @@ class ZioVertxServerTest extends TestSuite with OptionValues {
           basicRequest.get(baseUri).send(backend).map(_.body.toOption.value should not include "vert.x-eventloop-thread")
         }
       )
+      def drainZStream(zStream: ZioStreams.BinaryStream): Task[Unit] =
+        zStream.run(ZSink.drain)
 
       new AllServerTests(createServerTest, interpreter, backend, multipart = false, reject = false, options = false).tests() ++
         new ServerMultipartTests(
@@ -39,7 +42,7 @@ class ZioVertxServerTest extends TestSuite with OptionValues {
           partContentTypeHeaderSupport = false, // README: doesn't seem supported but I may be wrong
           partOtherHeaderSupport = false
         ).tests() ++ additionalTests() ++
-        new ServerStreamingTests(createServerTest, ZioStreams).tests() ++
+        new ServerStreamingTests(createServerTest, maxLengthSupported = true).tests(ZioStreams)(drainZStream) ++
         new ServerWebSocketTests(createServerTest, ZioStreams) {
           override def functionToPipe[A, B](f: A => B): streams.Pipe[A, B] = in => in.map(f)
           override def emptyPipe[A, B]: streams.Pipe[A, B] = _ => ZStream.empty