diff --git a/hscontrol/servertest/client.go b/hscontrol/servertest/client.go index 5a9ff5a4..54d87004 100644 --- a/hscontrol/servertest/client.go +++ b/hscontrol/servertest/client.go @@ -192,17 +192,39 @@ func (c *TestClient) register(tb testing.TB) { func (c *TestClient) startPoll(tb testing.TB) { tb.Helper() + c.startPollLoop() +} + +// startPollLoop creates a fresh poll context and launches the background +// [controlclient.Direct.PollNetMap] goroutine, which blocks until the +// context is cancelled or the server closes the connection. +func (c *TestClient) startPollLoop() { c.pollCtx, c.pollCancel = context.WithCancel(context.Background()) c.pollDone = make(chan struct{}) go func() { defer close(c.pollDone) - // [controlclient.Direct.PollNetMap] blocks until ctx is cancelled or the server closes - // the connection. _ = c.direct.PollNetMap(c.pollCtx, c) }() } +// resetNetmapState clears the cached netmap and drains any pending +// updates from a previous session so that convergence waits observe +// only the new session's maps. +func (c *TestClient) resetNetmapState() { + c.mu.Lock() + c.netmap = nil + c.mu.Unlock() + + for { + select { + case <-c.updates: + default: + return + } + } +} + // UpdateFullNetmap implements [controlclient.NetmapUpdater]. // Called by [controlclient.Direct] when a new [netmap.NetworkMap] is received. func (c *TestClient) UpdateFullNetmap(nm *netmap.NetworkMap) { @@ -274,24 +296,11 @@ func (c *TestClient) Reconnect(tb testing.TB) { } } - // Clear stale netmap data so that callers like [TestClient.WaitForPeers] - // actually wait for the new session's map instead of returning - // immediately based on the old session's cached state. - c.mu.Lock() - c.netmap = nil - c.mu.Unlock() - - // Drain any pending updates from the old session so they - // don't satisfy a subsequent [TestClient.WaitForPeers]/[TestClient.WaitForUpdate]. - for { - select { - case <-c.updates: - default: - goto drained - } - } - -drained: + // Clear stale netmap data and drain pending updates so that callers + // like [TestClient.WaitForPeers] actually wait for the new session's + // map instead of returning immediately based on the old session's + // cached state. + c.resetNetmapState() // Re-register and start polling again. c.register(tb) @@ -343,30 +352,8 @@ func (c *TestClient) ReloginAndPoll(ctx context.Context) error { return fmt.Errorf("servertest: TryLogin(%s): unexpected auth URL %q (expected auto-auth with preauth key)", c.Name, url) //nolint:err113 } - // Clear stale netmap state from the previous session so that - // convergence waits observe only the new session's maps. - c.mu.Lock() - c.netmap = nil - c.mu.Unlock() - - for { - select { - case <-c.updates: - continue - default: - } - - break - } - - c.pollCtx, c.pollCancel = context.WithCancel(context.Background()) - c.pollDone = make(chan struct{}) - - go func() { - defer close(c.pollDone) - - _ = c.direct.PollNetMap(c.pollCtx, c) - }() + c.resetNetmapState() + c.startPollLoop() return nil } @@ -389,14 +376,7 @@ func (c *TestClient) RestartPoll(ctx context.Context) error { } } - c.pollCtx, c.pollCancel = context.WithCancel(context.Background()) - c.pollDone = make(chan struct{}) - - go func() { - defer close(c.pollDone) - - _ = c.direct.PollNetMap(c.pollCtx, c) - }() + c.startPollLoop() return nil } @@ -430,10 +410,25 @@ func (c *TestClient) Netmap() *netmap.NetworkMap { func (c *TestClient) WaitForPeers(tb testing.TB, n int, timeout time.Duration) { tb.Helper() + c.waitForPeers(tb, n, timeout, "WaitForPeers", func(got int) bool { return got >= n }) +} + +// waitForPeers blocks until match reports the current peer count +// satisfies the caller's predicate, or until timeout expires. op +// names the caller for the timeout failure message. +func (c *TestClient) waitForPeers( + tb testing.TB, + n int, + timeout time.Duration, + op string, + match func(got int) bool, +) { + tb.Helper() + deadline := time.After(timeout) for { - if nm := c.Netmap(); nm != nil && len(nm.Peers) >= n { + if nm := c.Netmap(); nm != nil && match(len(nm.Peers)) { return } @@ -448,7 +443,7 @@ func (c *TestClient) WaitForPeers(tb testing.TB, n int, timeout time.Duration) { got = len(nm.Peers) } - tb.Fatalf("servertest: WaitForPeers(%s, %d): timeout after %v (got %d peers)", c.Name, n, timeout, got) + tb.Fatalf("servertest: %s(%s, %d): timeout after %v (got %d peers)", op, c.Name, n, timeout, got) } } } @@ -552,27 +547,7 @@ func (c *TestClient) SelfName() string { func (c *TestClient) WaitForPeerCount(tb testing.TB, n int, timeout time.Duration) { tb.Helper() - deadline := time.After(timeout) - - for { - if nm := c.Netmap(); nm != nil && len(nm.Peers) == n { - return - } - - select { - case <-c.updates: - // Check again. - case <-deadline: - nm := c.Netmap() - - got := 0 - if nm != nil { - got = len(nm.Peers) - } - - tb.Fatalf("servertest: WaitForPeerCount(%s, %d): timeout after %v (got %d peers)", c.Name, n, timeout, got) - } - } + c.waitForPeers(tb, n, timeout, "WaitForPeerCount", func(got int) bool { return got == n }) } // WaitForCondition blocks until condFn returns true on the latest diff --git a/hscontrol/servertest/harness.go b/hscontrol/servertest/harness.go index a88aa118..7461b1cb 100644 --- a/hscontrol/servertest/harness.go +++ b/hscontrol/servertest/harness.go @@ -76,8 +76,7 @@ func NewHarness(tb testing.TB, numClients int, opts ...HarnessOption) *TestHarne for i := range numClients { name := clientName(i) - copts := append([]ClientOption{WithUser(user)}, hc.clientOpts...) - c := NewClient(tb, server, name, copts...) + c := h.newClient(tb, name, hc.clientOpts...) h.clients = append(h.clients, c) } @@ -124,13 +123,22 @@ func (h *TestHarness) AddClient(tb testing.TB, opts ...ClientOption) *TestClient tb.Helper() name := clientName(len(h.clients)) - copts := append([]ClientOption{WithUser(h.defaultUser)}, opts...) - c := NewClient(tb, h.Server, name, copts...) + c := h.newClient(tb, name, opts...) h.clients = append(h.clients, c) return c } +// newClient creates a [TestClient] on the harness server, prepending the +// shared default user so caller options can override it. +func (h *TestHarness) newClient(tb testing.TB, name string, opts ...ClientOption) *TestClient { + tb.Helper() + + copts := append([]ClientOption{WithUser(h.defaultUser)}, opts...) + + return NewClient(tb, h.Server, name, copts...) +} + // WaitForMeshComplete blocks until every connected client sees // (connectedCount - 1) peers. func (h *TestHarness) WaitForMeshComplete(tb testing.TB, timeout time.Duration) { diff --git a/hscontrol/servertest/server.go b/hscontrol/servertest/server.go index fd20c519..71e3e6bb 100644 --- a/hscontrol/servertest/server.go +++ b/hscontrol/servertest/server.go @@ -236,40 +236,32 @@ func (s *TestServer) CreateUser(tb testing.TB, name string) *types.User { // CreatePreAuthKey creates a reusable pre-auth key for the given user. func (s *TestServer) CreatePreAuthKey(tb testing.TB, userID types.UserID) string { tb.Helper() - - uid := userID - - pak, err := s.st.CreatePreAuthKey(&uid, true, false, nil, nil) - if err != nil { - tb.Fatalf("servertest: CreatePreAuthKey: %v", err) - } - - return pak.Key + return s.createPreAuthKey(tb, userID, true, false, nil) } // CreateTaggedPreAuthKey creates a reusable pre-auth key with ACL tags. func (s *TestServer) CreateTaggedPreAuthKey(tb testing.TB, userID types.UserID, tags []string) string { tb.Helper() - - uid := userID - - pak, err := s.st.CreatePreAuthKey(&uid, true, false, nil, tags) - if err != nil { - tb.Fatalf("servertest: CreateTaggedPreAuthKey: %v", err) - } - - return pak.Key + return s.createPreAuthKey(tb, userID, true, false, tags) } // CreateEphemeralPreAuthKey creates an ephemeral pre-auth key. func (s *TestServer) CreateEphemeralPreAuthKey(tb testing.TB, userID types.UserID) string { tb.Helper() + return s.createPreAuthKey(tb, userID, false, true, nil) +} - uid := userID +func (s *TestServer) createPreAuthKey( + tb testing.TB, + userID types.UserID, + reusable, ephemeral bool, + tags []string, +) string { + tb.Helper() - pak, err := s.st.CreatePreAuthKey(&uid, false, true, nil, nil) + pak, err := s.st.CreatePreAuthKey(&userID, reusable, ephemeral, nil, tags) if err != nil { - tb.Fatalf("servertest: CreateEphemeralPreAuthKey: %v", err) + tb.Fatalf("servertest: createPreAuthKey: %v", err) } return pak.Key