rework test - still needs proper asserts

apache · Oct 30, 2023 · 2556190 · 2556190
1 parent 8914caa
commit 2556190
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/http2-tests/src/test/scala/org/apache/pekko/http/impl/engine/http2/Http2ServerSpec.scala b/http2-tests/src/test/scala/org/apache/pekko/http/impl/engine/http2/Http2ServerSpec.scala
@@ -1785,7 +1785,7 @@ class Http2ServerSpec extends PekkoSpecWithMaterializer("""
         dataStream.expectCancellation()
         terminated.futureValue
       })
-      "cancel connection during rapid reset attack".inAssertAllStagesStopped(new TestSetup with RequestResponseProbes {
+      "cancel connection during rapid reset attack".inAllStagesStopped(new TestSetup with RequestResponseProbes {
         implicit val bigEndian: ByteOrder = ByteOrder.BIG_ENDIAN
         val bb = new ByteStringBuilder
         bb.putInt(0)
@@ -1813,6 +1813,16 @@ class Http2ServerSpec extends PekkoSpecWithMaterializer("""
 
         // and then assert that all stages, substreams in particular, are stopped
       }
+
+    def inAllStagesStopped(runTest: => TestSetup) = {
+      val setup = runTest
+
+      // force connection to shutdown (in case it is an invalid state)
+      setup.network.fromNet.sendError(new RuntimeException)
+      setup.network.toNet.cancel()
+
+      // and then assert that all stages, substreams in particular, are stopped
+    }
   }
 
   protected /* To make ByteFlag warnings go away */ abstract class TestSetupWithoutHandshake {