stream assert

Author: rli

plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/encryption/JwtEncryptionManagerTest.kt

@@ -6,6 +6,7 @@ package software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.jupiter.api.Test
 import java.io.ByteArrayOutputStream
+import java.util.concurrent.atomic.AtomicBoolean
 import javax.crypto.spec.SecretKeySpec
 import kotlin.random.Random
 
@@ -40,18 +41,26 @@ class JwtEncryptionManagerTest {
         val bytes = "DEADBEEF".repeat(8).hexToByteArray() // 32 bytes
         val key = SecretKeySpec(bytes, "HmacSHA256")
 
-        val os = ByteArrayOutputStream()
+        val closed = AtomicBoolean(false)
+        val os = object : ByteArrayOutputStream() {
+            override fun close() {
+                closed.set(true)
+            }
+        }
         JwtEncryptionManager(key).writeInitializationPayload(os)
         assertThat(os.toString())
             // Flare requires encryption ends with new line
             // https://github.com/aws/language-server-runtimes/blob/4d7f81295dc12b59ed2e1c0ebaedb85ccb86cf76/runtimes/README.md#encryption
             .endsWith("\n")
-            // language=JSON
             .isEqualTo(
+                // language=JSON
                 """
             |{"version":"1.0","mode":"JWT","key":"3q2-796tvu_erb7v3q2-796tvu_erb7v3q2-796tvu8"}
             |
                 """.trimMargin()
             )
+
+        // writeInitializationPayload should not close the stream
+        assertThat(closed.get()).isFalse
     }
 }