Write response upstreamId in streaming (#419)

src/main/kotlin/io/emeraldpay/dshackle/rpc/NativeCall.kt

@@ -110,20 +110,34 @@ open class NativeCall(
         return if (callResult.stream == null) {
             Mono.just(buildResponse(callResult))
         } else {
-            val stream = callResult.stream.map { stream ->
+            return callResult.stream.switchOnFirst { t, stream ->
-                val result = BlockchainOuterClass.NativeCallReplyItem.newBuilder()
+                val firstChunk = t.get()
-                    .setSucceed(true)
+                if (firstChunk == null) {
-                    .setFinalChunk(stream.finalChunk)
+                    stream.map { buildStreamResult(it, callResult.id).build() }
-                    .setChunked(true)
+                } else {
-                    .setId(callResult.id)
+                    Flux.concat(
-                result.payload = ByteString.copyFrom(stream.chunkData)
+                        Mono.just(firstChunk)
-
+                            .map {
-                result.build()
+                                buildStreamResult(it, callResult.id)
+                                    .setUpstreamId(callResult.upstreamId)
+                                    .build()
+                            },
+                        stream.skip(1).map { buildStreamResult(it, callResult.id).build() },
+                    )
+                }
             }
-            stream
         }
     }
 
+    private fun buildStreamResult(chunk: Chunk, id: Int): BlockchainOuterClass.NativeCallReplyItem.Builder {
+        return BlockchainOuterClass.NativeCallReplyItem.newBuilder()
+            .setSucceed(true)
+            .setFinalChunk(chunk.finalChunk)
+            .setChunked(true)
+            .setPayload(ByteString.copyFrom(chunk.chunkData))
+            .setId(id)
+    }
+
     private fun completeSpan(callResult: CallResult, requestCount: Int) {
         val span = tracer.currentSpan()
         if (callResult.isError()) {

src/test/kotlin/io/emeraldpay/dshackle/rpc/NativeCallTest.kt

@@ -0,0 +1,132 @@
+package io.emeraldpay.dshackle.rpc
+
+import com.google.protobuf.ByteString
+import io.emeraldpay.api.proto.BlockchainOuterClass.NativeCallReplyItem
+import io.emeraldpay.api.proto.BlockchainOuterClass.NativeCallRequest
+import io.emeraldpay.dshackle.config.MainConfig
+import io.emeraldpay.dshackle.upstream.MultistreamHolder
+import io.emeraldpay.dshackle.upstream.rpcclient.stream.Chunk
+import io.emeraldpay.dshackle.upstream.signature.ResponseSigner
+import org.junit.jupiter.api.Test
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.spy
+import org.springframework.cloud.sleuth.Tracer
+import reactor.core.publisher.Flux
+import reactor.core.publisher.Mono
+import reactor.test.StepVerifier
+import java.time.Duration
+
+class NativeCallTest {
+
+    @Test
+    fun `nativeCall always returns item with response upstreamId`() {
+        val request = Mono.just(NativeCallRequest.newBuilder().build())
+        val nativeCall = spy(
+            NativeCall(
+                mock<MultistreamHolder>(),
+                mock<ResponseSigner>(),
+                MainConfig(),
+                mock<Tracer>(),
+            ),
+        ) {
+            on { nativeCallResult(request) } doReturn Flux.just(
+                NativeCall.CallResult.ok(1, null, "0x1".toByteArray(), null, "id", null),
+            )
+        }
+
+        StepVerifier.create(nativeCall.nativeCall(request))
+            .expectNext(
+                NativeCallReplyItem.newBuilder()
+                    .setUpstreamId("id")
+                    .setId(1)
+                    .setSucceed(true)
+                    .setPayload(ByteString.copyFrom("0x1".toByteArray()))
+                    .build(),
+            )
+            .expectComplete()
+            .verify(Duration.ofSeconds(1))
+    }
+
+    @Test
+    fun `nativeCall always returns error item with response upstreamId`() {
+        val request = Mono.just(NativeCallRequest.newBuilder().build())
+        val nativeCall = spy(
+            NativeCall(
+                mock<MultistreamHolder>(),
+                mock<ResponseSigner>(),
+                MainConfig(),
+                mock<Tracer>(),
+            ),
+        ) {
+            on { nativeCallResult(request) } doReturn Flux.just(
+                NativeCall.CallResult(1, null, null, NativeCall.CallError(50001, "message", null, null, "upId"), null, null),
+            )
+        }
+
+        StepVerifier.create(nativeCall.nativeCall(request))
+            .expectNext(
+                NativeCallReplyItem.newBuilder()
+                    .setUpstreamId("upId")
+                    .setId(1)
+                    .setSucceed(false)
+                    .setErrorMessage("message")
+                    .setItemErrorCode(50001)
+                    .build(),
+            )
+            .expectComplete()
+            .verify(Duration.ofSeconds(1))
+    }
+
+    @Test
+    fun `nativeCall always returns item with response upstreamId in the first chunk only`() {
+        val request = Mono.just(NativeCallRequest.newBuilder().build())
+        val chunks = Flux.just(
+            Chunk("0x1".toByteArray(), false),
+            Chunk("111".toByteArray(), false),
+            Chunk("222".toByteArray(), true),
+        )
+        val nativeCall = spy(
+            NativeCall(
+                mock<MultistreamHolder>(),
+                mock<ResponseSigner>(),
+                MainConfig(),
+                mock<Tracer>(),
+            ),
+        ) {
+            on { nativeCallResult(request) } doReturn Flux.just(
+                NativeCall.CallResult.ok(1, null, "".toByteArray(), null, "upId", null, chunks),
+            )
+        }
+
+        StepVerifier.create(nativeCall.nativeCall(request))
+            .expectNext(
+                NativeCallReplyItem.newBuilder()
+                    .setUpstreamId("upId")
+                    .setId(1)
+                    .setChunked(true)
+                    .setSucceed(true)
+                    .setPayload(ByteString.copyFrom("0x1".toByteArray()))
+                    .build(),
+            )
+            .expectNext(
+                NativeCallReplyItem.newBuilder()
+                    .setId(1)
+                    .setChunked(true)
+                    .setSucceed(true)
+                    .setPayload(ByteString.copyFrom("111".toByteArray()))
+                    .build(),
+            )
+            .expectNext(
+                NativeCallReplyItem.newBuilder()
+                    .setId(1)
+                    .setChunked(true)
+                    .setFinalChunk(true)
+                    .setSucceed(true)
+                    .setPayload(ByteString.copyFrom("222".toByteArray()))
+                    .build(),
+            )
+            .expectComplete()
+            .verify(Duration.ofSeconds(1))
+    }
+}