add test case for non-timeout cancellation

Author: halotukozak

rest/jetty/src/test/scala/io/udash/rest/jetty/CloseStaleJettyConnectionsOnMonixTimeout.scala

@@ -9,24 +9,29 @@ import monix.execution.atomic.Atomic
 import org.eclipse.jetty.client.HttpClient
 import org.eclipse.jetty.ee8.servlet.{ServletContextHandler, ServletHolder}
 import org.eclipse.jetty.server.{NetworkConnector, Server}
+import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AsyncFunSuite
 
 import java.net.InetSocketAddress
 import scala.concurrent.Future
 import scala.concurrent.duration.{FiniteDuration, IntMult}
 
-final class CloseStaleJettyConnectionsOnMonixTimeout extends AsyncFunSuite {
+final class CloseStaleJettyConnectionsOnMonixTimeout extends AsyncFunSuite with BeforeAndAfterEach {
 
-  test("close connection on monix task timeout") {
-    import monix.execution.Scheduler.Implicits.global
+  import monix.execution.Scheduler.Implicits.global
 
-    val MaxConnections: Int = 1 // to timeout quickly
-    val Connections: Int = 10 // > MaxConnections
-    val RequestTimeout: FiniteDuration = 1.hour // no timeout
-    val CallTimeout: FiniteDuration = 300.millis
+  private val MaxConnections: Int = 1 // to timeout quickly
+  private val Connections: Int = 10 // > MaxConnections
+  private val RequestTimeout: FiniteDuration = 1.hour // no timeout
+  private val CallTimeout: FiniteDuration = 300.millis
 
+  private var server: Server = _
+  private var httpClient: HttpClient = _
+  private var client: RestApiWithNeverCounter = _
 
-    val server = new Server(new InetSocketAddress("localhost", 0)) {
+  override def beforeEach(): Unit = {
+    super.beforeEach()
+    server = new Server(new InetSocketAddress("localhost", 0)) {
       setHandler(
         new ServletContextHandler().setup(
           _.addServlet(
@@ -40,32 +45,49 @@ final class CloseStaleJettyConnectionsOnMonixTimeout extends AsyncFunSuite {
       start()
     }
 
-    val httpClient = new HttpClient() {
+    httpClient = new HttpClient() {
       setMaxConnectionsPerDestination(MaxConnections)
       setIdleTimeout(RequestTimeout.toMillis)
       start()
     }
 
-    val client = JettyRestClient[RestApiWithNeverCounter](
+    client = JettyRestClient[RestApiWithNeverCounter](
       client = httpClient,
       baseUri = server.getConnectors.head |> { case connector: NetworkConnector => s"http://${connector.getHost}:${connector.getLocalPort}" },
       maxResponseLength = Int.MaxValue, // to avoid unnecessary logs
       timeout = RequestTimeout,
     )
+  }
 
+  override def afterEach(): Unit = {
+    RestApiWithNeverCounter.Impl.counter.set(0)
+    server.stop()
+    httpClient.stop()
+    super.afterEach()
+  }
+
+  test("close connection on monix task timeout") {
     Task
       .traverse(List.range(0, Connections))(_ => Task.fromFuture(client.neverGet).timeout(CallTimeout).failed)
       .timeoutTo(Connections * CallTimeout + 500.millis, Task(fail("All connections should have been closed"))) // + 500 millis just in case
       .map(_ => assert(RestApiWithNeverCounter.Impl.counter.get() == Connections)) // neverGet should be called Connections times
-      .guarantee(Task {
-        server.stop()
-        httpClient.stop()
-      })
+      .runToFuture
+  }
+
+  test("close connection on monix task cancellation") {
+    Task
+      .traverse(List.range(0, Connections)) { i =>
+        val cancelable = Task.fromFuture(client.neverGet).runAsync(_ => ())
+        Task.sleep(100.millis)
+          .restartUntil(_ => RestApiWithNeverCounter.Impl.counter.get() >= i)
+          .map(_ => cancelable.cancel())
+      }
+      .map(_ => assert(RestApiWithNeverCounter.Impl.counter.get() == Connections))
       .runToFuture
   }
 }
 
-object CloseStaleJettyConnectionsOnMonixTimeout {
+private object CloseStaleJettyConnectionsOnMonixTimeout {
   sealed trait RestApiWithNeverCounter {
     final val counter = Atomic(0)
     @GET def neverGet: Future[Unit]