netbird/client/internal/peer/wg_watcher_test.go

package peer

import (
	"context"
	"testing"
	"time"

	log "github.com/sirupsen/logrus"

	"github.com/netbirdio/netbird/client/iface/configurer"
)

type MocWgIface struct {
	initial       bool
	lastHandshake time.Time
	stop          bool
}

func (m *MocWgIface) GetStats(key string) (configurer.WGStats, error) {
	if !m.initial {
		m.initial = true
		return configurer.WGStats{}, nil
	}

	if !m.stop {
		m.lastHandshake = time.Now()
	}

	stats := configurer.WGStats{
		LastHandshake: m.lastHandshake,
	}

	return stats, nil
}

func (m *MocWgIface) disconnect() {
	m.stop = true
}

func TestWGWatcher_EnableWgWatcher(t *testing.T) {
	checkPeriod = 5 * time.Second
	wgHandshakeOvertime = 1 * time.Second

	mlog := log.WithField("peer", "tet")
	mocWgIface := &MocWgIface{}
	watcher := NewWGWatcher(mlog, mocWgIface, "")

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	onDisconnected := make(chan struct{}, 1)
	watcher.EnableWgWatcher(ctx, func() {
		mlog.Infof("onDisconnectedFn")
		onDisconnected <- struct{}{}
	})

	// wait for initial reading
	time.Sleep(2 * time.Second)
	mocWgIface.disconnect()

	select {
	case <-onDisconnected:
	case <-time.After(10 * time.Second):
		t.Errorf("timeout")
	}
	watcher.DisableWgWatcher()
}

func TestWGWatcher_ReEnable(t *testing.T) {
	checkPeriod = 5 * time.Second
	wgHandshakeOvertime = 1 * time.Second

	mlog := log.WithField("peer", "tet")
	mocWgIface := &MocWgIface{}
	watcher := NewWGWatcher(mlog, mocWgIface, "")

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	onDisconnected := make(chan struct{}, 1)

	watcher.EnableWgWatcher(ctx, func() {})
	watcher.DisableWgWatcher()

	watcher.EnableWgWatcher(ctx, func() {
		onDisconnected <- struct{}{}
	})

	time.Sleep(2 * time.Second)
	mocWgIface.disconnect()

	select {
	case <-onDisconnected:
	case <-time.After(10 * time.Second):
		t.Errorf("timeout")
	}
	watcher.DisableWgWatcher()
}