fix: merging Flows (#1032)

Author: lauzadis

aws-runtime/aws-core/api/aws-core.api

@@ -26,6 +26,9 @@ public class aws/sdk/kotlin/runtime/ConfigurationException : aws/sdk/kotlin/runt
 	public fun <init> (Ljava/lang/Throwable;)V
 }
 
+public final class aws/sdk/kotlin/runtime/FlowUtilKt {
+}
+
 public abstract interface annotation class aws/sdk/kotlin/runtime/InternalSdkApi : java/lang/annotation/Annotation {
 }
 

aws-runtime/aws-core/build.gradle.kts

@@ -8,6 +8,7 @@ extra["displayName"] = "AWS :: SDK :: Kotlin :: Client Runtime"
 extra["moduleName"] = "aws.sdk.kotlin.runtime"
 
 val smithyKotlinVersion: String by project
+val coroutinesVersion: String by project
 
 kotlin {
     sourceSets {
@@ -18,6 +19,12 @@ kotlin {
             }
         }
 
+        commonTest {
+            dependencies {
+                implementation("org.jetbrains.kotlinx:kotlinx-coroutines-test:$coroutinesVersion")
+            }
+        }
+
         all {
             languageSettings.optIn("aws.smithy.kotlin.runtime.InternalApi")
         }

aws-runtime/aws-core/common/src/aws/sdk/kotlin/runtime/FlowUtil.kt

@@ -0,0 +1,15 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package aws.sdk.kotlin.runtime
+
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.emitAll
+import kotlinx.coroutines.flow.flow
+
+// FIXME relocate to smithy-kotlin
+@InternalSdkApi
+public fun <T> mergeSequential(vararg flows: Flow<T>): Flow<T> = flow {
+    flows.forEach { flow -> emitAll(flow) }
+}

aws-runtime/aws-core/common/test/FlowUtilTest.kt

@@ -0,0 +1,58 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+import aws.sdk.kotlin.runtime.InternalSdkApi
+import aws.sdk.kotlin.runtime.mergeSequential
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+@OptIn(InternalSdkApi::class)
+class FlowUtilTest {
+
+    @Test
+    fun testMergingFlows() = runTest {
+        val a = flowOf(1)
+        val b = flowOf(2, 3, 4)
+        val merged = mergeSequential(a, b).toList()
+        assertEquals(listOf(1, 2, 3, 4), merged)
+    }
+
+    @Test
+    fun testMergingEmptyFlow() = runTest {
+        val a: Flow<Int> = flowOf()
+        val b: Flow<Int> = flowOf(4, 5, 6)
+
+        val merged = mergeSequential(a, b).toList()
+        assertEquals(listOf(4, 5, 6), merged)
+    }
+
+    @Test
+    fun testMergingOneFlow() = runTest {
+        val a = flowOf(1, 2, 3)
+        val merged = mergeSequential(a).toList()
+
+        assertEquals(listOf(1, 2, 3), merged)
+    }
+
+    @Test
+    fun testMergingSameFlow() = runTest {
+        val a = flowOf(1, 2, 3)
+        val merged = mergeSequential(a, a).toList()
+        assertEquals(listOf(1, 2, 3, 1, 2, 3), merged)
+    }
+
+    @Test
+    fun testMergingMoreThanTwoFlows() = runTest {
+        val a = flowOf(1)
+        val b = flowOf(2)
+        val c = flowOf(3)
+
+        val merged = mergeSequential(a, b, c).toList()
+        assertEquals(listOf(1, 2, 3), merged)
+    }
+}

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/AwsRuntimeTypes.kt

@@ -15,6 +15,7 @@ object AwsRuntimeTypes {
     object Core : RuntimeTypePackage(AwsKotlinDependency.AWS_CORE) {
         val AwsErrorMetadata = symbol("AwsErrorMetadata")
         val ClientException = symbol("ClientException")
+        val mergeSequential = symbol("mergeSequential")
 
         object Client : RuntimeTypePackage(AwsKotlinDependency.AWS_CORE, "client") {
             val AwsClientOption = symbol("AwsClientOption")

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/eventstream/EventStreamParserGenerator.kt

@@ -218,7 +218,7 @@ class EventStreamParserGenerator(
             .closeAndOpenBlock("} else {")
             .write(
                 "#T(#T(firstMessage), frames)",
-                RuntimeTypes.KotlinxCoroutines.Flow.merge,
+                AwsRuntimeTypes.Core.mergeSequential,
                 RuntimeTypes.KotlinxCoroutines.Flow.flowOf,
             )
             .closeBlock("}")

codegen/smithy-aws-kotlin-codegen/src/main/kotlin/aws/sdk/kotlin/codegen/protocols/eventstream/EventStreamSerializerGenerator.kt

@@ -5,6 +5,7 @@
 
 package aws.sdk.kotlin.codegen.protocols.eventstream
 
+import aws.sdk.kotlin.codegen.AwsRuntimeTypes
 import software.amazon.smithy.codegen.core.CodegenException
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.kotlin.codegen.core.*
@@ -93,7 +94,7 @@ class EventStreamSerializerGenerator(
             writer.withBlock(
                 "val messages = #T(#T(initialRequest), stream.#T(::#T))",
                 "",
-                RuntimeTypes.KotlinxCoroutines.Flow.merge,
+                AwsRuntimeTypes.Core.mergeSequential,
                 RuntimeTypes.KotlinxCoroutines.Flow.flowOf,
                 RuntimeTypes.KotlinxCoroutines.Flow.map,
                 encodeFn,