110 lines
3.4 KiB
Go
110 lines
3.4 KiB
Go
// 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)
|
|
})
|
|
}
|