servertest: bound fixture batching and await observable state

Updates #3544

hscontrol/servertest/consistency_test.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/juanfont/headscale/hscontrol/servertest"
 	"github.com/stretchr/testify/assert"
+	"tailscale.com/types/netmap"
 )
 
 // TestConsistency verifies that all nodes converge to the same
@@ -86,14 +87,27 @@ func TestConsistency(t *testing.T) {
 		c6 := h.AddClient(t)
 		c7 := h.AddClient(t)
 
-		// Wait for new nodes to see at least all other connected
-		// clients (they may also see the disconnected nodes during
-		// the grace period, so we check >= not ==).
+		// Disconnected peers can satisfy a count before the new peers arrive.
+		// Wait for each connected identity before comparing the mesh.
 		connected := h.ConnectedClients()
-		minPeers := len(connected) - 1
-
 		for _, c := range connected {
-			c.WaitForPeers(t, minPeers, 30*time.Second)
+			c.WaitForCondition(t, "all connected peers visible", 30*time.Second,
+				func(nm *netmap.NetworkMap) bool {
+					seen := make(map[string]bool, len(nm.Peers))
+					for _, peer := range nm.Peers {
+						if hi := peer.Hostinfo(); hi.Valid() {
+							seen[hi.Hostname()] = true
+						}
+					}
+
+					for _, other := range connected {
+						if other != c && !seen[other.Name] {
+							return false
+						}
+					}
+
+					return true
+				})
 		}
 
 		// Verify the new nodes can see each other.

hscontrol/servertest/grants_test.go

@@ -224,12 +224,28 @@ func TestGrantPolicies(t *testing.T) { //nolint:gocyclo
 		c1.WaitForCondition(t, "first grant policy update",
 			10*time.Second,
 			func(nm *netmap.NetworkMap) bool {
-				return c1.UpdateCount() > countC1
+				for _, rule := range nm.PacketFilter {
+					for _, dst := range rule.Dsts {
+						if dst.Ports.First == 22 && dst.Ports.Last == 22 {
+							return c1.UpdateCount() > countC1
+						}
+					}
+				}
+
+				return false
 			})
 		c2.WaitForCondition(t, "first grant policy update",
 			10*time.Second,
 			func(nm *netmap.NetworkMap) bool {
-				return c2.UpdateCount() > countC2
+				for _, rule := range nm.PacketFilter {
+					for _, dst := range rule.Dsts {
+						if dst.Ports.First == 22 && dst.Ports.Last == 22 {
+							return c2.UpdateCount() > countC2
+						}
+					}
+				}
+
+				return false
 			})
 
 		// Capture the update count after first policy.
@@ -270,12 +286,12 @@ func TestGrantPolicies(t *testing.T) { //nolint:gocyclo
 		c1.WaitForCondition(t, "second grant policy update",
 			10*time.Second,
 			func(nm *netmap.NetworkMap) bool {
-				return c1.UpdateCount() > updateCountPhase1
+				return c1.UpdateCount() > updateCountPhase1 && hasCapMatches(nm.PacketFilter)
 			})
 		c2.WaitForCondition(t, "second grant policy update",
 			10*time.Second,
 			func(nm *netmap.NetworkMap) bool {
-				return c2.UpdateCount() > countC2
+				return c2.UpdateCount() > countC2 && hasCapMatches(nm.PacketFilter)
 			})
 
 		// Verify the second policy added cap grants that weren't
@@ -751,25 +767,27 @@ func TestGrantViaSubnetFilterRules(t *testing.T) {
 	// Critical: the router's [netmap.NetworkMap.PacketFilter] MUST contain rules with
 	// the via-steered subnet (10.0.0.0/24) as a destination.
 	// Without this, the router drops traffic forwarded through it.
-	routerNM := routerA.Netmap()
-	require.NotNil(t, routerNM)
-	require.NotNil(t, routerNM.PacketFilter,
-		"router PacketFilter should not be nil")
+	requireNetmapHolds(t, func(tt require.TestingT) {
+		routerNM := routerA.Netmap()
+		require.NotNil(tt, routerNM)
+		require.NotNil(tt, routerNM.PacketFilter,
+			"router PacketFilter should not be nil")
 
-	var foundSubnetDst bool
+		var foundSubnetDst bool
 
-	for _, m := range routerNM.PacketFilter {
-		for _, dst := range m.Dsts {
-			dstPrefix := netip.PrefixFrom(dst.Net.Addr(), dst.Net.Bits())
-			if route.Contains(dstPrefix.Addr()) && dstPrefix.Bits() >= route.Bits() {
-				foundSubnetDst = true
+		for _, m := range routerNM.PacketFilter {
+			for _, dst := range m.Dsts {
+				dstPrefix := netip.PrefixFrom(dst.Net.Addr(), dst.Net.Bits())
+				if route.Contains(dstPrefix.Addr()) && dstPrefix.Bits() >= route.Bits() {
+					foundSubnetDst = true
+				}
 			}
 		}
-	}
 
-	assert.True(t, foundSubnetDst,
-		"router PacketFilter should contain destination rules for via-steered subnet 10.0.0.0/24; "+
-			"without per-node filter compilation for via grants, these rules are missing")
+		assert.True(tt, foundSubnetDst,
+			"router PacketFilter should contain destination rules for via-steered subnet 10.0.0.0/24; "+
+				"without per-node filter compilation for via grants, these rules are missing")
+	})
 }
 
 // TestGrantViaSubnetBroaderDstFilterRules verifies that a via grant
@@ -868,25 +886,27 @@ func TestGrantViaSubnetBroaderDstFilterRules(t *testing.T) {
 
 	// Router's PacketFilter must contain the broader grant dst as a
 	// destination — that is what Tailscale SaaS emits for via grants.
-	routerNM := routerA.Netmap()
-	require.NotNil(t, routerNM)
-	require.NotNil(t, routerNM.PacketFilter,
-		"router PacketFilter should not be nil")
+	requireNetmapHolds(t, func(tt require.TestingT) {
+		routerNM := routerA.Netmap()
+		require.NotNil(tt, routerNM)
+		require.NotNil(tt, routerNM.PacketFilter,
+			"router PacketFilter should not be nil")
 
-	var foundBroaderDst bool
+		var foundBroaderDst bool
 
-	for _, m := range routerNM.PacketFilter {
-		for _, dst := range m.Dsts {
-			dstPrefix := netip.PrefixFrom(dst.Net.Addr(), dst.Net.Bits())
-			if dstPrefix == broaderDst {
-				foundBroaderDst = true
+		for _, m := range routerNM.PacketFilter {
+			for _, dst := range m.Dsts {
+				dstPrefix := netip.PrefixFrom(dst.Net.Addr(), dst.Net.Bits())
+				if dstPrefix == broaderDst {
+					foundBroaderDst = true
+				}
 			}
 		}
-	}
 
-	assert.True(t, foundBroaderDst,
-		"router PacketFilter should contain destination rules for the broader grant dst 10.0.0.0/8; "+
-			"the via gate requires advertised-route overlap, and the emitted prefix is the dst")
+		assert.True(tt, foundBroaderDst,
+			"router PacketFilter should contain destination rules for the broader grant dst 10.0.0.0/8; "+
+				"the via gate requires advertised-route overlap, and the emitted prefix is the dst")
+	})
 }
 
 // TestGrantViaExitNodeNoFilterRules verifies wire-format SaaS compat:

hscontrol/servertest/policy_test.go

@@ -110,20 +110,28 @@ func TestPolicyChanges(t *testing.T) {
 		t.Parallel()
 		h := servertest.NewHarness(t, 2)
 
-		// Apply policy twice and verify updates arrive both times.
-		for round := range 2 {
+		// Change the rule each time; reapplying a policy is a no-op.
+		for round, port := range []uint16{22, 23} {
 			countBefore := h.Client(0).UpdateCount()
 
-			h.ChangePolicy(t, []byte(`{
+			h.ChangePolicy(t, fmt.Appendf(nil, `{
 				"acls": [
-					{"action": "accept", "src": ["*"], "dst": ["*:*"]}
+					{"action": "accept", "src": ["*"], "dst": ["*:%d"]}
 				]
-			}`))
+			}`, port))
 
 			h.Client(0).WaitForCondition(t, "update after policy change",
 				10*time.Second,
 				func(nm *netmap.NetworkMap) bool {
-					return h.Client(0).UpdateCount() > countBefore
+					for _, rule := range nm.PacketFilter {
+						for _, dst := range rule.Dsts {
+							if dst.Ports.First == port && dst.Ports.Last == port {
+								return h.Client(0).UpdateCount() > countBefore
+							}
+						}
+					}
+
+					return false
 				})
 
 			t.Logf("round %d: update received", round)

hscontrol/servertest/race_test.go

@@ -385,9 +385,18 @@ func TestRaceConnectDuringGracePeriod(t *testing.T) {
 		surviving := originals[n/2:]
 		allActive := append(surviving, replacements...)
 
-		for _, c := range allActive {
-			c.WaitForPeers(t, len(allActive)-1, 30*time.Second)
-		}
+		requireNetmapHolds(t, func(tt require.TestingT) {
+			for _, client := range allActive {
+				for _, peer := range allActive {
+					if client == peer {
+						continue
+					}
+
+					_, found := client.PeerByName(peer.Name)
+					assert.True(tt, found, "%s must see active peer %s", client.Name, peer.Name)
+				}
+			}
+		})
 
 		servertest.AssertConsistentState(t, allActive)
 	})

hscontrol/servertest/server.go

@@ -146,6 +146,8 @@ func NewServer(tb testing.TB, opts ...ServerOption) *TestServer {
 		},
 		Taildrop: types.TaildropConfig{Enabled: sc.taildropEnabled},
 		Tuning: types.Tuning{
+			// Small fixtures rarely fill a batch; bound their write latency.
+			NodeStoreBatchTimeout:          10 * time.Millisecond,
 			BatchChangeDelay:               sc.batchDelay,
 			BatcherWorkers:                 sc.batcherWorkers,
 			NodeMapSessionBufferedChanSize: sc.bufferedChanSize,

hscontrol/servertest/stress_test.go

@@ -624,10 +624,20 @@ func TestStressChurn(t *testing.T) {
 				servertest.WithUser(user))
 		}
 
-		// Wait for the new set to converge.
-		for _, c := range clients {
-			c.WaitForPeers(t, n-1, 30*time.Second)
-		}
+		// Old disconnected peers can satisfy the count before replacements
+		// arrive. Wait for every replacement identity to stay visible.
+		requireNetmapHolds(t, func(tt require.TestingT) {
+			for _, client := range clients {
+				for _, peer := range clients {
+					if client == peer {
+						continue
+					}
+
+					_, found := client.PeerByName(peer.Name)
+					assert.True(tt, found, "%s must see replacement %s", client.Name, peer.Name)
+				}
+			}
+		})
 
 		servertest.AssertSymmetricVisibility(t, clients)
 	})