snapshot of https://github.com/livekit/livekit.git @ 08c5433449c106943231037d0b892895f6e92d3b (2026-09-15)

pkg/service/rtcservice_test.go

@@ -0,0 +1,109 @@
+// Copyright 2026 LiveKit, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package service
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/proto"
+
+	"github.com/livekit/protocol/livekit"
+
+	"github.com/livekit/livekit-server/pkg/routing"
+)
+
+func leaveResponse(action livekit.LeaveRequest_Action) *livekit.SignalResponse {
+	return &livekit.SignalResponse{
+		Message: &livekit.SignalResponse_Leave{
+			Leave: &livekit.LeaveRequest{
+				Action: action,
+				Reason: livekit.DisconnectReason_MIGRATION,
+			},
+		},
+	}
+}
+
+func TestDrainMessageSource(t *testing.T) {
+	collect := func(msgs *[]proto.Message) func(proto.Message) bool {
+		return func(msg proto.Message) bool {
+			*msgs = append(*msgs, msg)
+			return true
+		}
+	}
+
+	t.Run("drains queued messages of a closed source", func(t *testing.T) {
+		source := routing.NewDefaultMessageChannel("CO_test")
+		require.NoError(t, source.WriteMessage(leaveResponse(livekit.LeaveRequest_RESUME)))
+		require.NoError(t, source.WriteMessage(leaveResponse(livekit.LeaveRequest_RECONNECT)))
+		source.Close()
+
+		var got []proto.Message
+		require.True(t, drainMessageSource(source, time.Second, collect(&got)))
+		require.Len(t, got, 2)
+		require.Equal(
+			t,
+			livekit.LeaveRequest_RESUME,
+			got[0].(*livekit.SignalResponse).GetLeave().GetAction(),
+		)
+		require.Equal(
+			t,
+			livekit.LeaveRequest_RECONNECT,
+			got[1].(*livekit.SignalResponse).GetLeave().GetAction(),
+		)
+	})
+
+	t.Run("drains messages written while draining", func(t *testing.T) {
+		source := routing.NewDefaultMessageChannel("CO_test")
+		// mimics the relay pushing a message that was still in flight when the
+		// request direction went away, and closing the source right after
+		go func() {
+			time.Sleep(20 * time.Millisecond)
+			_ = source.WriteMessage(leaveResponse(livekit.LeaveRequest_RESUME))
+			source.Close()
+		}()
+
+		var got []proto.Message
+		require.True(t, drainMessageSource(source, time.Second, collect(&got)))
+		require.Len(t, got, 1)
+	})
+
+	t.Run("gives up on the deadline when the source stays open", func(t *testing.T) {
+		source := routing.NewDefaultMessageChannel("CO_test")
+		require.NoError(t, source.WriteMessage(leaveResponse(livekit.LeaveRequest_RESUME)))
+
+		var got []proto.Message
+		start := time.Now()
+		require.False(t, drainMessageSource(source, 50*time.Millisecond, collect(&got)))
+		require.GreaterOrEqual(t, time.Since(start), 50*time.Millisecond)
+		// what was queued is still flushed
+		require.Len(t, got, 1)
+	})
+
+	t.Run("stops when the write fails", func(t *testing.T) {
+		source := routing.NewDefaultMessageChannel("CO_test")
+		require.NoError(t, source.WriteMessage(leaveResponse(livekit.LeaveRequest_RESUME)))
+		require.NoError(t, source.WriteMessage(leaveResponse(livekit.LeaveRequest_RECONNECT)))
+		source.Close()
+
+		var got []proto.Message
+		require.False(t, drainMessageSource(source, time.Second, func(msg proto.Message) bool {
+			got = append(got, msg)
+			return false
+		}))
+		require.Len(t, got, 1)
+	})
+}