Add unit tests for posture checks and peer location resolution logic

This commit is contained in:
mlsmaycon
2026-06-23 14:32:22 +02:00
parent b4c1db17e4
commit 330a03ce75
3 changed files with 269 additions and 1 deletions

View File

@@ -49,6 +49,7 @@ import (
nbdns "github.com/netbirdio/netbird/dns"
"github.com/netbirdio/netbird/management/server/activity"
"github.com/netbirdio/netbird/management/server/geolocation"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
"github.com/netbirdio/netbird/management/server/posture"
"github.com/netbirdio/netbird/management/server/store"
@@ -2893,3 +2894,141 @@ func TestUpdatePeer_DnsLabelUniqueName(t *testing.T) {
require.NoError(t, err, "renaming to unique FQDN should succeed")
assert.Equal(t, "api-server", updated.DNSLabel, "DNS label should be first label of FQDN")
}
// fakeGeo is a configurable geolocation.Geolocation implementation for tests. It
// returns a record built from the configured city geoname id, or an error when set.
type fakeGeo struct {
geoNameID uint
isoCode string
cityName string
err error
}
func (g *fakeGeo) Lookup(net.IP) (*geolocation.Record, error) {
if g.err != nil {
return nil, g.err
}
record := &geolocation.Record{}
record.City.GeonameID = g.geoNameID
record.City.Names.En = g.cityName
record.Country.ISOCode = g.isoCode
return record, nil
}
func (g *fakeGeo) GetAllCountries() ([]geolocation.Country, error) { return nil, nil }
func (g *fakeGeo) GetCitiesByCountry(string) ([]geolocation.City, error) { return nil, nil }
func (g *fakeGeo) Stop() error { return nil }
func TestResolvePeerLocation(t *testing.T) {
realIP := net.ParseIP("203.0.113.10")
tests := []struct {
name string
geo geolocation.Geolocation
peer *nbpeer.Peer
realIP net.IP
want *nbpeer.Location
wantNil bool
}{
{
name: "no geo configured returns nil",
geo: nil,
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
wantNil: true,
},
{
name: "nil real IP returns nil",
geo: &fakeGeo{geoNameID: 100},
peer: &nbpeer.Peer{ID: "p1"},
realIP: nil,
wantNil: true,
},
{
name: "lookup error returns nil",
geo: &fakeGeo{err: fmt.Errorf("lookup boom")},
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
wantNil: true,
},
{
name: "same IP and same geoname returns nil",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: realIP,
GeoNameID: 100,
},
},
realIP: realIP,
wantNil: true,
},
{
name: "same IP but changed geoname returns location",
geo: &fakeGeo{geoNameID: 200, isoCode: "US", cityName: "City B"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: realIP,
GeoNameID: 100,
},
},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City B",
GeoNameID: 200,
},
},
{
name: "different IP returns location",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{
ID: "p1",
Location: nbpeer.Location{
ConnectionIP: net.ParseIP("198.51.100.7"),
GeoNameID: 100,
},
},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City A",
GeoNameID: 100,
},
},
{
name: "no prior location returns location",
geo: &fakeGeo{geoNameID: 100, isoCode: "US", cityName: "City A"},
peer: &nbpeer.Peer{ID: "p1"},
realIP: realIP,
want: &nbpeer.Location{
ConnectionIP: realIP,
CountryCode: "US",
CityName: "City A",
GeoNameID: 100,
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
am := &DefaultAccountManager{geo: tt.geo}
got := am.resolvePeerLocation(context.Background(), tt.peer, tt.realIP)
if tt.wantNil {
assert.Nil(t, got, "resolved location should be nil")
return
}
require.NotNil(t, got, "resolved location should not be nil")
assert.True(t, tt.want.ConnectionIP.Equal(got.ConnectionIP), "connection IP should match")
assert.Equal(t, tt.want.CountryCode, got.CountryCode, "country code should match")
assert.Equal(t, tt.want.CityName, got.CityName, "city name should match")
assert.Equal(t, tt.want.GeoNameID, got.GeoNameID, "geoname id should match")
})
}
}

View File

@@ -0,0 +1,129 @@
package posture
import (
"net/netip"
"testing"
"github.com/stretchr/testify/assert"
nbpeer "github.com/netbirdio/netbird/management/server/peer"
)
func TestAffectsPosture(t *testing.T) {
processCheck := &Checks{Checks: ChecksDefinition{ProcessCheck: &ProcessCheck{}}}
osCheck := &Checks{Checks: ChecksDefinition{OSVersionCheck: &OSVersionCheck{}}}
nbCheck := &Checks{Checks: ChecksDefinition{NBVersionCheck: &NBVersionCheck{}}}
geoCheck := &Checks{Checks: ChecksDefinition{GeoLocationCheck: &GeoLocationCheck{}}}
privateRangeCheck := &Checks{Checks: ChecksDefinition{
PeerNetworkRangeCheck: &PeerNetworkRangeCheck{
Ranges: []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")},
},
}}
publicRangeCheck := &Checks{Checks: ChecksDefinition{
PeerNetworkRangeCheck: &PeerNetworkRangeCheck{
Ranges: []netip.Prefix{netip.MustParsePrefix("203.0.113.0/24")},
},
}}
mixedRangeCheck := &Checks{Checks: ChecksDefinition{
PeerNetworkRangeCheck: &PeerNetworkRangeCheck{
Ranges: []netip.Prefix{
netip.MustParsePrefix("203.0.113.0/24"),
netip.MustParsePrefix("192.168.0.0/16"),
},
},
}}
tests := []struct {
name string
diff *nbpeer.MetaDiff
checks []*Checks
want bool
}{
{
name: "nil diff never affects posture",
diff: nil,
checks: []*Checks{processCheck},
want: false,
},
{
name: "process check affected by files change",
diff: &nbpeer.MetaDiff{Files: true},
checks: []*Checks{processCheck},
want: true,
},
{
name: "process check ignores unrelated change",
diff: &nbpeer.MetaDiff{Hostname: true},
checks: []*Checks{processCheck},
want: false,
},
{
name: "os check affected by os version change",
diff: &nbpeer.MetaDiff{OSVersion: true},
checks: []*Checks{osCheck},
want: true,
},
{
name: "nb check affected by wt version change",
diff: &nbpeer.MetaDiff{WtVersion: true},
checks: []*Checks{nbCheck},
want: true,
},
{
name: "geo check affected by location change",
diff: &nbpeer.MetaDiff{LocationChanged: true},
checks: []*Checks{geoCheck},
want: true,
},
{
name: "network range check not affected without network address or location change",
diff: &nbpeer.MetaDiff{Hostname: true},
checks: []*Checks{privateRangeCheck},
want: false,
},
{
name: "private range check affected by network address change",
diff: &nbpeer.MetaDiff{NetworkAddresses: true},
checks: []*Checks{privateRangeCheck},
want: true,
},
{
name: "public range check not affected by network address change alone",
diff: &nbpeer.MetaDiff{NetworkAddresses: true},
checks: []*Checks{publicRangeCheck},
want: false,
},
{
name: "public range check affected by location change alone",
diff: &nbpeer.MetaDiff{LocationChanged: true},
checks: []*Checks{publicRangeCheck},
want: true,
},
{
name: "private range check affected by location change alone",
diff: &nbpeer.MetaDiff{LocationChanged: true},
checks: []*Checks{privateRangeCheck},
want: true,
},
{
name: "public range check affected when location also changed",
diff: &nbpeer.MetaDiff{NetworkAddresses: true, LocationChanged: true},
checks: []*Checks{publicRangeCheck},
want: true,
},
{
name: "mixed ranges affected by network address change due to private range",
diff: &nbpeer.MetaDiff{NetworkAddresses: true},
checks: []*Checks{mixedRangeCheck},
want: true,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := AffectsPosture(tt.diff, tt.checks)
assert.Equal(t, tt.want, got, "AffectsPosture result should match expectation")
})
}
}

View File

@@ -73,7 +73,7 @@ func AffectsPosture(diff *nbpeer.MetaDiff, checks []*Checks) bool {
if c.Checks.GeoLocationCheck != nil && diff.LocationChanged {
return true
}
if c.Checks.PeerNetworkRangeCheck != nil && diff.NetworkAddresses {
if c.Checks.PeerNetworkRangeCheck != nil && (diff.NetworkAddresses || diff.LocationChanged) {
if diff.LocationChanged {
return true
}