mirror of
https://github.com/netbirdio/netbird.git
synced 2025-08-06 14:10:01 +02:00
0dab03252c
- Clients now subscribe to peer status changes. - The server manages and maintains these subscriptions. - Replaced raw string peer IDs with a custom peer ID type for better type safety and clarity.
100 lines
2.5 KiB
Go
100 lines
2.5 KiB
Go
package client
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/netbirdio/netbird/relay/messages"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type mockRelayedConn struct {
|
|
}
|
|
|
|
func (m *mockRelayedConn) Write(p []byte) (n int, err error) {
|
|
return len(p), nil
|
|
}
|
|
|
|
func TestWaitToBeOnlineAndSubscribe_Success(t *testing.T) {
|
|
peerID := messages.HashID("peer1")
|
|
mockConn := &mockRelayedConn{}
|
|
logger := logrus.New()
|
|
logger.SetOutput(&bytes.Buffer{}) // discard log output
|
|
sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
// Launch wait in background
|
|
go func() {
|
|
time.Sleep(100 * time.Millisecond)
|
|
sub.OnPeersOnline([]messages.PeerID{peerID})
|
|
}()
|
|
|
|
err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
|
|
assert.NoError(t, err)
|
|
}
|
|
|
|
func TestWaitToBeOnlineAndSubscribe_Timeout(t *testing.T) {
|
|
peerID := messages.HashID("peer2")
|
|
mockConn := &mockRelayedConn{}
|
|
logger := logrus.New()
|
|
logger.SetOutput(&bytes.Buffer{})
|
|
sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
|
|
defer cancel()
|
|
|
|
err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
|
|
assert.Error(t, err)
|
|
assert.Equal(t, context.DeadlineExceeded, err)
|
|
}
|
|
|
|
func TestWaitToBeOnlineAndSubscribe_Duplicate(t *testing.T) {
|
|
peerID := messages.HashID("peer3")
|
|
mockConn := &mockRelayedConn{}
|
|
logger := logrus.New()
|
|
logger.SetOutput(&bytes.Buffer{})
|
|
sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
|
|
|
|
ctx := context.Background()
|
|
go func() {
|
|
_ = sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
|
|
|
|
}()
|
|
time.Sleep(100 * time.Millisecond)
|
|
err := sub.WaitToBeOnlineAndSubscribe(ctx, peerID)
|
|
require.Error(t, err)
|
|
assert.Contains(t, err.Error(), "already waiting")
|
|
}
|
|
|
|
func TestUnsubscribeStateChange(t *testing.T) {
|
|
peerID := messages.HashID("peer4")
|
|
mockConn := &mockRelayedConn{}
|
|
logger := logrus.New()
|
|
logger.SetOutput(&bytes.Buffer{})
|
|
sub := NewPeersStateSubscription(logrus.NewEntry(logger), mockConn, nil)
|
|
|
|
doneChan := make(chan struct{})
|
|
go func() {
|
|
_ = sub.WaitToBeOnlineAndSubscribe(context.Background(), peerID)
|
|
close(doneChan)
|
|
}()
|
|
time.Sleep(100 * time.Millisecond)
|
|
|
|
err := sub.UnsubscribeStateChange([]messages.PeerID{peerID})
|
|
assert.NoError(t, err)
|
|
|
|
select {
|
|
case <-doneChan:
|
|
case <-time.After(200 * time.Millisecond):
|
|
// Expected timeout, meaning the subscription was successfully unsubscribed
|
|
t.Errorf("timeout")
|
|
}
|
|
}
|