netbird/proxy/handle_mapping_stream_test.go

package proxy

import (
	"context"
	"io"
	"testing"

	log "github.com/sirupsen/logrus"
	"github.com/stretchr/testify/assert"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"

	"github.com/netbirdio/netbird/proxy/internal/health"
	"github.com/netbirdio/netbird/shared/management/proto"
)

type mockMappingStream struct {
	grpc.ClientStream
	messages []*proto.GetMappingUpdateResponse
	idx      int
}

func (m *mockMappingStream) Recv() (*proto.GetMappingUpdateResponse, error) {
	if m.idx >= len(m.messages) {
		return nil, io.EOF
	}
	msg := m.messages[m.idx]
	m.idx++
	return msg, nil
}

func (m *mockMappingStream) Header() (metadata.MD, error) {
	return nil, nil //nolint:nilnil
}
func (m *mockMappingStream) Trailer() metadata.MD     { return nil }
func (m *mockMappingStream) CloseSend() error         { return nil }
func (m *mockMappingStream) Context() context.Context { return context.Background() }
func (m *mockMappingStream) SendMsg(any) error        { return nil }
func (m *mockMappingStream) RecvMsg(any) error        { return nil }

func closedChan() chan struct{} {
	ch := make(chan struct{})
	close(ch)
	return ch
}

func TestHandleMappingStream_SyncCompleteFlag(t *testing.T) {
	checker := health.NewChecker(nil, nil)
	s := &Server{
		Logger:        log.StandardLogger(),
		healthChecker: checker,
		routerReady:   closedChan(),
	}

	stream := &mockMappingStream{
		messages: []*proto.GetMappingUpdateResponse{
			{InitialSyncComplete: true},
		},
	}

	syncDone := false
	err := s.handleMappingStream(context.Background(), stream, &syncDone)
	assert.NoError(t, err)
	assert.True(t, syncDone, "initial sync should be marked done when flag is set")
}

func TestHandleMappingStream_NoSyncFlagDoesNotMarkDone(t *testing.T) {
	checker := health.NewChecker(nil, nil)
	s := &Server{
		Logger:        log.StandardLogger(),
		healthChecker: checker,
		routerReady:   closedChan(),
	}

	stream := &mockMappingStream{
		messages: []*proto.GetMappingUpdateResponse{
			{}, // no sync flag
		},
	}

	syncDone := false
	err := s.handleMappingStream(context.Background(), stream, &syncDone)
	assert.NoError(t, err)
	assert.False(t, syncDone, "initial sync should not be marked done without flag")
}

func TestHandleMappingStream_NilHealthChecker(t *testing.T) {
	s := &Server{
		Logger:      log.StandardLogger(),
		routerReady: closedChan(),
	}

	stream := &mockMappingStream{
		messages: []*proto.GetMappingUpdateResponse{
			{InitialSyncComplete: true},
		},
	}

	syncDone := false
	err := s.handleMappingStream(context.Background(), stream, &syncDone)
	assert.NoError(t, err)
	assert.True(t, syncDone, "sync done flag should be set even without health checker")
}