From c1e8f2823495bae59bf7fb52dd198a291ac65768 Mon Sep 17 00:00:00 2001
From: Oleg Bespalov <oleg.bespalov@grafana.com>
Date: Mon, 11 Sep 2023 09:20:40 +0200
Subject: [PATCH] Dedicated unit test

This adds a dedicated unit test for checking that all responses from the
server were received even if client declared that sending is finished
(client.end) called.
---
 grpc/stream_test.go | 96 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/grpc/stream_test.go b/grpc/stream_test.go
index fb8ba69..7b0e928 100644
--- a/grpc/stream_test.go
+++ b/grpc/stream_test.go
@@ -3,6 +3,7 @@ package grpc
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/dop251/goja"
 	"github.com/grafana/xk6-grpc/grpc/testutils/grpcservice"
@@ -188,6 +189,101 @@ func TestStream_ErrorHandling(t *testing.T) {
 	)
 }
 
+// this test case is checking that everything that server sends
+// after the client finished (client.end called) is delivered to the client
+// and the end event is called
+func TestStream_ReceiveAllServerResponsesAfterEnd(t *testing.T) {
+	t.Parallel()
+
+	ts := newTestState(t)
+
+	stub := &FeatureExplorerStub{}
+
+	savedFeatures := []*grpcservice.Feature{
+		{
+			Name: "foo",
+			Location: &grpcservice.Point{
+				Latitude:  1,
+				Longitude: 2,
+			},
+		},
+		{
+			Name: "bar",
+			Location: &grpcservice.Point{
+				Latitude:  3,
+				Longitude: 4,
+			},
+		},
+	}
+
+	stub.listFeatures = func(rect *grpcservice.Rectangle, stream grpcservice.FeatureExplorer_ListFeaturesServer) error {
+		for _, feature := range savedFeatures {
+			// adding a delay to make server response "slower"
+			time.Sleep(200 * time.Millisecond)
+
+			if err := stream.Send(feature); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	}
+
+	grpcservice.RegisterFeatureExplorerServer(ts.httpBin.ServerGRPC, stub)
+
+	replace := func(code string) (goja.Value, error) {
+		return ts.VU.Runtime().RunString(ts.httpBin.Replacer.Replace(code))
+	}
+
+	initString := codeBlock{
+		code: `
+		var client = new grpc.Client();
+		client.load([], "../grpc/testutils/grpcservice/route_guide.proto");`,
+	}
+	vuString := codeBlock{
+		code: `
+		client.connect("GRPCBIN_ADDR");
+		let stream = new grpc.Stream(client, "main.FeatureExplorer/ListFeatures")
+		stream.on('data', function (data) {
+			call('Feature:' + data.name);
+		});
+		stream.on('end', function () {
+			call('End called');
+		});
+
+		stream.write({
+			lo: {
+			  latitude: 1,
+			  longitude: 2,
+			},
+			hi: {
+			  latitude: 1,
+			  longitude: 2,
+			},
+		});
+		stream.end();
+		`,
+	}
+
+	val, err := replace(initString.code)
+	assertResponse(t, initString, err, val, ts)
+
+	ts.ToVUContext()
+
+	val, err = replace(vuString.code)
+
+	ts.EventLoop.WaitOnRegistered()
+
+	assertResponse(t, vuString, err, val, ts)
+
+	assert.Equal(t, ts.callRecorder.Recorded(), []string{
+		"Feature:foo",
+		"Feature:bar",
+		"End called",
+	},
+	)
+}
+
 // FeatureExplorerStub is a stub for FeatureExplorerServer
 // it has ability to override methods
 type FeatureExplorerStub struct {