mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-24 23:59:08 +02:00
94 lines
2.8 KiB
Go
94 lines
2.8 KiB
Go
package grpc
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"go.uber.org/mock/gomock"
|
|
|
|
"github.com/netbirdio/netbird/management/internals/modules/reverseproxy/proxy"
|
|
rpservice "github.com/netbirdio/netbird/management/internals/modules/reverseproxy/service"
|
|
"github.com/netbirdio/netbird/shared/management/proto"
|
|
)
|
|
|
|
const (
|
|
versionTestProxyID = "proxy-a"
|
|
versionTestCluster = "cluster.example.com"
|
|
versionTestVersion = "0.60.0"
|
|
)
|
|
|
|
// hangupStream cancels its context on the first Send, emulating a proxy that
|
|
// disconnects right after receiving the initial snapshot. The legacy stream
|
|
// carries no proxy-to-management messages, so this is the only way for
|
|
// GetMappingUpdate to return.
|
|
type hangupStream struct {
|
|
recordingStream
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
func (s *hangupStream) Send(m *proto.GetMappingUpdateResponse) error {
|
|
s.cancel()
|
|
return s.recordingStream.Send(m)
|
|
}
|
|
|
|
func (s *hangupStream) Context() context.Context { return s.ctx }
|
|
|
|
// newVersionTestServer wires a server whose proxy manager only accepts a
|
|
// Connect carrying versionTestVersion, so a dropped or mangled version fails
|
|
// the test as an unexpected call.
|
|
func newVersionTestServer(t *testing.T) *ProxyServiceServer {
|
|
t.Helper()
|
|
ctrl := gomock.NewController(t)
|
|
|
|
svcMgr := rpservice.NewMockManager(ctrl)
|
|
svcMgr.EXPECT().GetGlobalServices(gomock.Any()).Return(nil, nil)
|
|
|
|
proxyMgr := proxy.NewMockManager(ctrl)
|
|
proxyMgr.EXPECT().
|
|
Connect(gomock.Any(), versionTestProxyID, gomock.Any(), versionTestCluster, gomock.Any(), versionTestVersion, gomock.Any(), gomock.Any()).
|
|
Return(&proxy.Proxy{ID: versionTestProxyID, Version: versionTestVersion}, nil)
|
|
proxyMgr.EXPECT().Disconnect(gomock.Any(), versionTestProxyID, gomock.Any()).Return(nil)
|
|
|
|
s := newSnapshotTestServer(t, 10)
|
|
s.serviceManager = svcMgr
|
|
s.proxyManager = proxyMgr
|
|
return s
|
|
}
|
|
|
|
func TestSyncMappings_ForwardsProxyVersion(t *testing.T) {
|
|
s := newVersionTestServer(t)
|
|
|
|
// The init carries the version, the ack acknowledges the empty snapshot,
|
|
// and the exhausted fake stream then ends the RPC.
|
|
stream := &syncRecordingStream{
|
|
recvMsgs: []*proto.SyncMappingsRequest{
|
|
{Msg: &proto.SyncMappingsRequest_Init{Init: &proto.SyncMappingsInit{
|
|
ProxyId: versionTestProxyID,
|
|
Address: versionTestCluster,
|
|
Version: versionTestVersion,
|
|
}}},
|
|
ackMsg(),
|
|
},
|
|
}
|
|
|
|
err := s.SyncMappings(stream)
|
|
require.ErrorContains(t, err, "no more recv messages")
|
|
}
|
|
|
|
func TestGetMappingUpdate_ForwardsProxyVersion(t *testing.T) {
|
|
s := newVersionTestServer(t)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
t.Cleanup(cancel)
|
|
stream := &hangupStream{ctx: ctx, cancel: cancel}
|
|
|
|
err := s.GetMappingUpdate(&proto.GetMappingUpdateRequest{
|
|
ProxyId: versionTestProxyID,
|
|
Address: versionTestCluster,
|
|
Version: versionTestVersion,
|
|
}, stream)
|
|
require.ErrorIs(t, err, context.Canceled)
|
|
}
|