From bb151fbf77ecceecaff19412cabbe1ea574b204b Mon Sep 17 00:00:00 2001
From: Adam Warski <adam@warski.org>
Date: Tue, 6 Aug 2024 10:16:51 +0200
Subject: [PATCH] Add test for reading headers from web socket responses
 (#2245)

---
 .../client4/testing/websocket/WebSocketTest.scala  | 14 ++++++++++++++
 .../sttp/client4/FetchBackendWebSocketTest.scala   |  1 +
 .../client4/HttpClientFutureWebSocketTest.scala    |  2 ++
 .../sttp/client4/HttpClientSyncWebSocketTest.scala |  2 ++
 .../client4/impl/cats/FetchCatsWebSocketTest.scala |  1 +
 .../client4/impl/cats/FetchCatsWebSocketTest.scala |  1 +
 .../cats/HttpClientCatsWebSocketTest.scala         |  2 ++
 .../fs2/HttpClientFs2WebSocketTest.scala           |  2 ++
 .../fs2/HttpClientFs2WebSocketTest.scala           |  2 ++
 .../impl/monix/FetchMonixWebSocketTest.scala       |  1 +
 .../impl/monix/HttpClientMonixWebSocketTest.scala  |  2 ++
 .../client4/impl/zio/FetchZioWebSocketTest.scala   |  1 +
 .../zio/HttpClientZioWebSocketTest.scala           |  2 ++
 .../client4/impl/zio/FetchZioWebSocketTest.scala   |  1 +
 .../zio/HttpClientZioWebSocketTest.scala           |  2 ++
 .../sttp/client4/testing/server/HttpServer.scala   | 10 ++++++++++
 16 files changed, 46 insertions(+)

diff --git a/core/src/test/scala/sttp/client4/testing/websocket/WebSocketTest.scala b/core/src/test/scala/sttp/client4/testing/websocket/WebSocketTest.scala
index 44705a3cbe..7786a253f6 100644
--- a/core/src/test/scala/sttp/client4/testing/websocket/WebSocketTest.scala
+++ b/core/src/test/scala/sttp/client4/testing/websocket/WebSocketTest.scala
@@ -29,6 +29,7 @@ abstract class WebSocketTest[F[_]]
   implicit def monad: MonadError[F]
 
   def throwsWhenNotAWebSocket: Boolean = false
+  def supportsReadingWebSocketResponseHeaders: Boolean = true
 
   it should "send and receive three messages using asWebSocketAlways" in {
     basicRequest
@@ -207,6 +208,19 @@ abstract class WebSocketTest[F[_]]
       .toFuture()
   }
 
+  if (supportsReadingWebSocketResponseHeaders) {
+    it should "receive the extra headers set by the server" in {
+      basicRequest
+        .get(uri"$wsEndpoint/ws/header")
+        .response(asWebSocketAlways((ws: WebSocket[F]) => ws.close()))
+        .send(backend)
+        .map { response =>
+          response.header("Correlation-id") shouldBe Some("ABC-XYZ-123")
+        }
+        .toFuture()
+    }
+  }
+
   def sendText(ws: WebSocket[F], count: Int): F[Unit] =
     send(ws, count, (i: Int) => WebSocketFrame.text(s"test$i"))
 
diff --git a/core/src/test/scalajs/sttp/client4/FetchBackendWebSocketTest.scala b/core/src/test/scalajs/sttp/client4/FetchBackendWebSocketTest.scala
index 3988fbf67a..0962042045 100644
--- a/core/src/test/scalajs/sttp/client4/FetchBackendWebSocketTest.scala
+++ b/core/src/test/scalajs/sttp/client4/FetchBackendWebSocketTest.scala
@@ -12,6 +12,7 @@ class FetchBackendWebSocketTest extends WebSocketTest[Future] {
 
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketBackend[Future] = FetchBackend()
   override implicit val convertToFuture: ConvertToFuture[Future] = ConvertToFuture.future
diff --git a/core/src/test/scalajvm/sttp/client4/HttpClientFutureWebSocketTest.scala b/core/src/test/scalajvm/sttp/client4/HttpClientFutureWebSocketTest.scala
index d531e29fb7..1abdd7efc9 100644
--- a/core/src/test/scalajvm/sttp/client4/HttpClientFutureWebSocketTest.scala
+++ b/core/src/test/scalajvm/sttp/client4/HttpClientFutureWebSocketTest.scala
@@ -13,4 +13,6 @@ class HttpClientFutureWebSocketTest[F[_]] extends WebSocketTest[Future] with Web
   override implicit val monad: MonadError[Future] = new FutureMonad()
 
   override def concurrently[T](fs: List[() => Future[T]]): Future[List[T]] = Future.sequence(fs.map(_()))
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/core/src/test/scalajvm/sttp/client4/HttpClientSyncWebSocketTest.scala b/core/src/test/scalajvm/sttp/client4/HttpClientSyncWebSocketTest.scala
index 1a528296fa..ee91ad464d 100644
--- a/core/src/test/scalajvm/sttp/client4/HttpClientSyncWebSocketTest.scala
+++ b/core/src/test/scalajvm/sttp/client4/HttpClientSyncWebSocketTest.scala
@@ -12,4 +12,6 @@ class HttpClientSyncWebSocketTest extends WebSocketTest[Identity] {
   override implicit val monad: MonadError[Identity] = IdentityMonad
 
   override def throwsWhenNotAWebSocket: Boolean = true
+  // HttpClient doesn't expose the response headers for web sockets in any way
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/cats-ce2/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala b/effects/cats-ce2/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
index 31f2e806ec..39cbe765cc 100644
--- a/effects/cats-ce2/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
+++ b/effects/cats-ce2/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
@@ -10,6 +10,7 @@ import scala.scalajs.concurrent.JSExecutionContext.queue
 class FetchCatsWebSocketTest extends WebSocketTest[IO] with CatsTestBase {
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketBackend[IO] = FetchCatsBackend()
 }
diff --git a/effects/cats/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala b/effects/cats/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
index 31f2e806ec..39cbe765cc 100644
--- a/effects/cats/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
+++ b/effects/cats/src/test/scalajs/sttp/client4/impl/cats/FetchCatsWebSocketTest.scala
@@ -10,6 +10,7 @@ import scala.scalajs.concurrent.JSExecutionContext.queue
 class FetchCatsWebSocketTest extends WebSocketTest[IO] with CatsTestBase {
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketBackend[IO] = FetchCatsBackend()
 }
diff --git a/effects/cats/src/test/scalajvm/sttp/client4/httpclient/cats/HttpClientCatsWebSocketTest.scala b/effects/cats/src/test/scalajvm/sttp/client4/httpclient/cats/HttpClientCatsWebSocketTest.scala
index b71e1c9c7d..134b933f2b 100644
--- a/effects/cats/src/test/scalajvm/sttp/client4/httpclient/cats/HttpClientCatsWebSocketTest.scala
+++ b/effects/cats/src/test/scalajvm/sttp/client4/httpclient/cats/HttpClientCatsWebSocketTest.scala
@@ -10,4 +10,6 @@ class HttpClientCatsWebSocketTest
     with HttpClientCatsTestBase {
 
   override def concurrently[T](fs: List[() => IO[T]]): IO[List[T]] = fs.map(_()).parSequence
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/fs2-ce2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala b/effects/fs2-ce2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
index 5fdf46c056..ce617c8b6b 100644
--- a/effects/fs2-ce2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
+++ b/effects/fs2-ce2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
@@ -30,4 +30,6 @@ class HttpClientFs2WebSocketTest
     to.andThen(rest => fs2.Stream.eval(item.pure[IO]) ++ rest)
 
   override def concurrently[T](fs: List[() => IO[T]]): IO[List[T]] = fs.map(_()).parSequence
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/fs2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala b/effects/fs2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
index 5fdf46c056..ce617c8b6b 100644
--- a/effects/fs2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
+++ b/effects/fs2/src/test/scalajvm/sttp/client4/httpclient/fs2/HttpClientFs2WebSocketTest.scala
@@ -30,4 +30,6 @@ class HttpClientFs2WebSocketTest
     to.andThen(rest => fs2.Stream.eval(item.pure[IO]) ++ rest)
 
   override def concurrently[T](fs: List[() => IO[T]]): IO[List[T]] = fs.map(_()).parSequence
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/monix/src/test/scalajs/sttp/client4/impl/monix/FetchMonixWebSocketTest.scala b/effects/monix/src/test/scalajs/sttp/client4/impl/monix/FetchMonixWebSocketTest.scala
index c1c31ee658..bd7da50dd4 100644
--- a/effects/monix/src/test/scalajs/sttp/client4/impl/monix/FetchMonixWebSocketTest.scala
+++ b/effects/monix/src/test/scalajs/sttp/client4/impl/monix/FetchMonixWebSocketTest.scala
@@ -15,6 +15,7 @@ import scala.scalajs.concurrent.JSExecutionContext.Implicits.queue
 class FetchMonixWebSocketTest extends WebSocketTest[Task] with WebSocketStreamingTest[Task, MonixStreams] {
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketStreamBackend[Task, MonixStreams] = FetchMonixBackend()
   override implicit val convertToFuture: ConvertToFuture[Task] = convertMonixTaskToFuture
diff --git a/effects/monix/src/test/scalajvm/sttp/client4/impl/monix/HttpClientMonixWebSocketTest.scala b/effects/monix/src/test/scalajvm/sttp/client4/impl/monix/HttpClientMonixWebSocketTest.scala
index 003284c45c..651b7552eb 100644
--- a/effects/monix/src/test/scalajvm/sttp/client4/impl/monix/HttpClientMonixWebSocketTest.scala
+++ b/effects/monix/src/test/scalajvm/sttp/client4/impl/monix/HttpClientMonixWebSocketTest.scala
@@ -36,4 +36,6 @@ class HttpClientMonixWebSocketTest
     to.andThen(rest => Observable.now(item) ++ rest)
 
   override def concurrently[T](fs: List[() => Task[T]]): Task[List[T]] = Task.parSequence(fs.map(_()))
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/zio/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala b/effects/zio/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
index 810c765ca2..6242dc9237 100644
--- a/effects/zio/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
+++ b/effects/zio/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
@@ -15,6 +15,7 @@ import scala.scalajs.concurrent.JSExecutionContext.Implicits.queue
 class FetchZioWebSocketTest extends WebSocketTest[Task] with WebSocketStreamingTest[Task, ZioStreams] with ZioTestBase {
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketStreamBackend[Task, ZioStreams] = FetchZioBackend()
   override implicit val convertToFuture: ConvertToFuture[Task] = convertZioTaskToFuture
diff --git a/effects/zio/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala b/effects/zio/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
index 520ca5ba53..d02106b236 100644
--- a/effects/zio/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
+++ b/effects/zio/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
@@ -37,4 +37,6 @@ class HttpClientZioWebSocketTest
     to.andThen(rest => ZStream(item) ++ rest)
 
   override def concurrently[T](fs: List[() => Task[T]]): Task[List[T]] = ZIO.collectAllPar(fs.map(_()))
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/effects/zio1/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala b/effects/zio1/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
index 810c765ca2..6242dc9237 100644
--- a/effects/zio1/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
+++ b/effects/zio1/src/test/scalajs/sttp/client4/impl/zio/FetchZioWebSocketTest.scala
@@ -15,6 +15,7 @@ import scala.scalajs.concurrent.JSExecutionContext.Implicits.queue
 class FetchZioWebSocketTest extends WebSocketTest[Task] with WebSocketStreamingTest[Task, ZioStreams] with ZioTestBase {
   implicit override def executionContext: ExecutionContext = queue
   override def throwsWhenNotAWebSocket: Boolean = true
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 
   override val backend: WebSocketStreamBackend[Task, ZioStreams] = FetchZioBackend()
   override implicit val convertToFuture: ConvertToFuture[Task] = convertZioTaskToFuture
diff --git a/effects/zio1/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala b/effects/zio1/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
index ec90878786..a301a36e93 100644
--- a/effects/zio1/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
+++ b/effects/zio1/src/test/scalajvm/sttp/client4/httpclient/zio/HttpClientZioWebSocketTest.scala
@@ -37,4 +37,6 @@ class HttpClientZioWebSocketTest
     to.andThen(rest => ZStream(item) ++ rest)
 
   override def concurrently[T](fs: List[() => Task[T]]): Task[List[T]] = Task.collectAllPar(fs.map(_()))
+
+  override def supportsReadingWebSocketResponseHeaders: Boolean = false
 }
diff --git a/testing/server/src/main/scala/sttp/client4/testing/server/HttpServer.scala b/testing/server/src/main/scala/sttp/client4/testing/server/HttpServer.scala
index b9b09ecdab..518bcb0d42 100644
--- a/testing/server/src/main/scala/sttp/client4/testing/server/HttpServer.scala
+++ b/testing/server/src/main/scala/sttp/client4/testing/server/HttpServer.scala
@@ -4,6 +4,7 @@ import akka.actor.ActorSystem
 import akka.http.scaladsl.Http
 import akka.http.scaladsl.coding.Coders._
 import akka.http.scaladsl.coding.DeflateNoWrap
+import akka.http.scaladsl.model.HttpHeader.ParsingResult
 import akka.http.scaladsl.model._
 import akka.http.scaladsl.model.headers.CacheDirectives._
 import akka.http.scaladsl.model.headers._
@@ -487,6 +488,15 @@ private class HttpServer(port: Int, info: String => Unit) extends AutoCloseable
               )
             )
           }
+        } ~
+        path("header") {
+          respondWithHeader(HttpHeader.parse("Correlation-ID", "ABC-XYZ-123").asInstanceOf[ParsingResult.Ok].header) {
+            handleWebSocketMessages(Flow[Message].mapConcat {
+              case tm: TextMessage =>
+                TextMessage(Source.single("echo: ") ++ tm.textStream) :: Nil
+              case bm: BinaryMessage => bm :: Nil
+            })
+          }
         }
     } ~ path("empty_content_encoding") {
       get {