Fix reliability findings

clients/go/mxgateway/client_session_test.go

@@ -113,6 +113,40 @@ func TestEventSubscriptionCloseStopsStream(t *testing.T) {
 	}
 }
 
+func TestEventsAfterCancelsStreamWhenCompatibilityChannelIsAbandoned(t *testing.T) {
+	fake := &fakeGatewayServer{
+		streamStarted:    make(chan struct{}),
+		streamDone:       make(chan struct{}),
+		streamEventCount: 64,
+	}
+	client, cleanup := newBufconnClient(t, fake)
+	defer cleanup()
+	session := NewSessionForID(client, "session-1")
+
+	events, err := session.EventsAfter(context.Background(), 0)
+	if err != nil {
+		t.Fatalf("EventsAfter() error = %v", err)
+	}
+	<-fake.streamStarted
+
+	select {
+	case <-fake.streamDone:
+	case <-time.After(2 * time.Second):
+		t.Fatal("compatibility event stream did not stop after result channel filled")
+	}
+
+	for {
+		select {
+		case _, ok := <-events:
+			if !ok {
+				return
+			}
+		case <-time.After(2 * time.Second):
+			t.Fatal("compatibility event channel did not close")
+		}
+	}
+}
+
 func TestSessionHelpersBuildCommandsAndExposeRawReply(t *testing.T) {
 	fake := &fakeGatewayServer{
 		invokeReply: &pb.MxCommandReply{
@@ -267,13 +301,14 @@ func newBufconnClient(t *testing.T, fake *fakeGatewayServer) (*Client, func()) {
 type fakeGatewayServer struct {
 	pb.UnimplementedMxAccessGatewayServer
 
-	openReply     *pb.OpenSessionReply
-	openAuth      string
-	streamAuth    string
-	streamStarted chan struct{}
-	streamDone    chan struct{}
-	invokeReply   *pb.MxCommandReply
-	invokeRequest *pb.MxCommandRequest
+	openReply        *pb.OpenSessionReply
+	openAuth         string
+	streamAuth       string
+	streamStarted    chan struct{}
+	streamDone       chan struct{}
+	streamEventCount int
+	invokeReply      *pb.MxCommandReply
+	invokeRequest    *pb.MxCommandRequest
 }
 
 func (s *fakeGatewayServer) OpenSession(ctx context.Context, req *pb.OpenSessionRequest) (*pb.OpenSessionReply, error) {
@@ -320,12 +355,18 @@ func (s *fakeGatewayServer) StreamEvents(req *pb.StreamEventsRequest, stream grp
 	if s.streamStarted != nil {
 		close(s.streamStarted)
 	}
-	if err := stream.Send(&pb.MxEvent{
-		SessionId:      req.GetSessionId(),
-		Family:         pb.MxEventFamily_MX_EVENT_FAMILY_ON_DATA_CHANGE,
-		WorkerSequence: 1,
-	}); err != nil {
-		return err
+	eventCount := s.streamEventCount
+	if eventCount == 0 {
+		eventCount = 1
+	}
+	for sequence := 1; sequence <= eventCount; sequence++ {
+		if err := stream.Send(&pb.MxEvent{
+			SessionId:      req.GetSessionId(),
+			Family:         pb.MxEventFamily_MX_EVENT_FAMILY_ON_DATA_CHANGE,
+			WorkerSequence: uint64(sequence),
+		}); err != nil {
+			return err
+		}
 	}
 	<-stream.Context().Done()
 	return io.EOF