routewatch/internal/streamer/streamer_internal_test.go

package streamer

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"io"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"runtime"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/logger"
	"git.eeqj.de/sneak/routewatch/internal/metrics"
	"git.eeqj.de/sneak/routewatch/internal/ristypes"
)

func TestNewStreamer(t *testing.T) {
	t.Parallel()

	logger := logger.New()
	metricsTracker := metrics.New()
	s := New(logger, metricsTracker)

	if s == nil {
		t.Fatal("New() returned nil")
	}

	if s.logger != logger {
		t.Error("logger not set correctly")
	}

	if s.client == nil {
		t.Error("HTTP client not initialized")
	}

	if s.handlers == nil {
		t.Error("handlers slice not initialized")
	}

	if s.metrics != metricsTracker {
		t.Error("metrics tracker not set correctly")
	}

	if s.idleTimeout != streamIdleTimeout {
		t.Errorf("idleTimeout = %s, want %s", s.idleTimeout, streamIdleTimeout)
	}

	if s.minBackoff != minBackoffDelay {
		t.Errorf("minBackoff = %s, want %s", s.minBackoff, minBackoffDelay)
	}
}

// TestStreamDoesNotLeakTickersAcrossReconnects drives many short-lived
// connections (each stream call is one reconnect cycle) and asserts the
// goroutine count returns to its starting value. Each connection starts two
// ticker goroutines; before the fix they lived until the streamer's lifetime
// context was cancelled, so every reconnect leaked two.
//
//nolint:paralleltest // counts goroutines, so no other test may run alongside
func TestStreamDoesNotLeakTickersAcrossReconnects(t *testing.T) {
	// The handler returns immediately, so the response body is empty and each
	// stream call ends at once, standing in for a dropped connection.
	srv := httptest.NewServer(http.HandlerFunc(
		func(_ http.ResponseWriter, _ *http.Request) {}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL

	// One warm-up connection so any persistent HTTP transport goroutine exists
	// before we take the baseline.
	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("warm-up stream returned %v, want %v", err, errStreamEnded)
	}

	s.client.CloseIdleConnections()

	baseline := settledGoroutineCount()

	const reconnects = 20
	for range reconnects {
		err := s.stream(context.Background())
		if !errors.Is(err, errStreamEnded) {
			t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
		}
	}

	s.client.CloseIdleConnections()

	if !waitForGoroutines(baseline) {
		t.Fatalf("goroutines did not return to baseline %d after %d reconnects, got %d",
			baseline, reconnects, runtime.NumGoroutine())
	}
}

// updateLine is one UPDATE message as RIS Live sends it.
const updateLine = `{"type":"ris_message","data":{"type":"UPDATE",` +
	`"peer":"192.0.2.1","peer_asn":"64496","timestamp":1700000000}}` + "\n"

// updateHandler wants UPDATE messages. For each one it handles it sends on
// handled if there is room, and otherwise does nothing; a nil handled is
// never sent on.
type updateHandler struct {
	handled chan struct{}
}

func (updateHandler) WantsMessage(messageType string) bool {
	return messageType == "UPDATE"
}

func (updateHandler) QueueCapacity() int { return 10 }

func (h updateHandler) HandleMessage(*ristypes.RISMessage) {
	select {
	case h.handled <- struct{}{}:
	default:
	}
}

// flush sends what w has buffered to the client. The writers httptest servers
// hand to handlers can always do that.
func flush(w http.ResponseWriter) {
	flusher, ok := w.(http.Flusher)
	if ok {
		flusher.Flush()
	}
}

// stateDownLine is the STATE message RIS Live sends when the collector's BGP
// session with the peer stateDownPeer goes down.
const (
	stateDownPeer = "192.0.2.1"
	stateDownLine = `{"type":"ris_message","data":{"type":"STATE",` +
		`"peer":"` + stateDownPeer + `",` +
		`"peer_asn":"64496","timestamp":1700000000,"state":"down"}}` + "\n"
)

// stateHandler wants STATE messages only. Tests read what it is handed from
// its queue, so HandleMessage does nothing.
type stateHandler struct{}

func (stateHandler) WantsMessage(messageType string) bool {
	return messageType == ristypes.MessageTypeState
}

func (stateHandler) QueueCapacity() int { return 10 }

func (stateHandler) HandleMessage(*ristypes.RISMessage) {}

// TestStreamHandsStateMessagesToHandlersThatWantThem has the server send one
// STATE message whose state is down. The streamer must queue it, with its
// peer and state, for a handler that wants STATE messages.
func TestStreamHandsStateMessagesToHandlersThatWantThem(t *testing.T) {
	t.Parallel()

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, stateDownLine)
		}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL
	s.RegisterHandler(stateHandler{})

	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
	}

	select {
	case msg := <-s.handlers[0].queue:
		if msg.Type != ristypes.MessageTypeState || msg.Peer != stateDownPeer ||
			msg.State != ristypes.StateDown {
			t.Errorf("queued message has type %q, peer %q, state %q; "+
				"want %q, %q, %q", msg.Type, msg.Peer, msg.State,
				ristypes.MessageTypeState, stateDownPeer, ristypes.StateDown)
		}
	default:
		t.Fatal("the STATE message was not queued for the handler that wants it")
	}
}

// TestHeldStateMessagesKeepTheirPlaceInTheQueue hands queueMessage two STATE
// messages and two updates while the queue is full, making room in between.
// Neither STATE message may be dropped: each must reach the queue after
// everything queued before it and before anything handed over after it. The
// update handed over while a STATE message is held is dropped.
func TestHeldStateMessagesKeepTheirPlaceInTheQueue(t *testing.T) {
	t.Parallel()

	s := New(logger.New(), metrics.New())
	s.RegisterHandler(stateHandler{})

	info := s.handlers[0]

	before := &ristypes.RISMessage{Type: ristypes.MessageTypeUpdate}
	for len(info.queue) < cap(info.queue) {
		info.queue <- before
	}

	firstDown := &ristypes.RISMessage{
		Type: ristypes.MessageTypeState, State: ristypes.StateDown,
	}
	secondDown := &ristypes.RISMessage{
		Type: ristypes.MessageTypeState, State: ristypes.StateDown,
	}
	updateWhileHeld := &ristypes.RISMessage{Type: ristypes.MessageTypeUpdate}
	after := &ristypes.RISMessage{Type: ristypes.MessageTypeUpdate}

	if !s.queueMessage(info, firstDown) {
		t.Fatal("the first STATE message was dropped on a full queue")
	}

	<-info.queue // the handler makes room for one message

	if !s.queueMessage(info, secondDown) {
		t.Fatal("the second STATE message was dropped")
	}

	if s.queueMessage(info, updateWhileHeld) {
		t.Fatal("an update was queued while a STATE message was held")
	}

	// The handler takes every message queued before the first STATE message.
	for range cap(info.queue) - 1 {
		if <-info.queue != before {
			t.Fatal("a message overtook those queued before the first STATE message")
		}
	}

	if !s.queueMessage(info, after) {
		t.Fatal("an update was dropped on a queue with room")
	}

	if <-info.queue != firstDown {
		t.Fatal("the first STATE message did not come next")
	}

	if <-info.queue != secondDown {
		t.Fatal("the second STATE message did not come after the first")
	}

	if <-info.queue != after {
		t.Fatal("the update handed over after the STATE messages did not come " +
			"after them")
	}

	if len(info.queue) != 0 {
		t.Errorf("%d messages left on the queue, want none", len(info.queue))
	}
}

// takeAll takes the messages off queue, as a handler's worker would, and sends
// them on the returned channel once queue is closed.
func takeAll(queue chan *ristypes.RISMessage) chan []*ristypes.RISMessage {
	taken := make(chan []*ristypes.RISMessage, 1)

	go func() {
		var msgs []*ristypes.RISMessage

		for msg := range queue {
			msgs = append(msgs, msg)
		}

		taken <- msgs
	}()

	return taken
}

// newStreamerWithFullStateQueue returns a streamer reading from url with one
// handler, a stateHandler whose queue is full.
func newStreamerWithFullStateQueue(url string) *Streamer {
	s := New(logger.New(), metrics.New())
	s.url = url
	s.RegisterHandler(stateHandler{})

	queue := s.handlers[0].queue
	for len(queue) < cap(queue) {
		queue <- &ristypes.RISMessage{Type: ristypes.MessageTypeState}
	}

	return s
}

// TestOtherHandlersKeepReceivingWhileAStateMessageIsHeld has the server send a
// STATE message whose state is down, then an UPDATE, while the queue of the
// handler that wants STATE messages is full. The STATE message must be held
// for that handler, and the read loop must go on to queue the UPDATE for the
// handler that wants it.
func TestOtherHandlersKeepReceivingWhileAStateMessageIsHeld(t *testing.T) {
	t.Parallel()

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, stateDownLine+updateLine)
		}))
	defer srv.Close()

	s := newStreamerWithFullStateQueue(srv.URL)
	s.RegisterHandler(updateHandler{})

	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
	}

	held := s.handlers[0].heldState
	if len(held) != 1 || held[0].Peer != stateDownPeer ||
		held[0].State != ristypes.StateDown {
		t.Errorf("held %d STATE messages, want the down from %s alone",
			len(held), stateDownPeer)
	}

	got := len(s.handlers[1].queue)
	if got != 1 {
		t.Errorf("%d messages queued for the handler that wants UPDATE "+
			"messages, want 1", got)
	}
}

// TestStopQueuesTheHeldStateMessages has the server send two STATE messages
// and an UPDATE and hold the response open while the queue of the handler that
// wants STATE messages is full, so both are held. Stop must put both on the
// queue, after the messages already on it, before it closes the queue, and
// must end the stream.
func TestStopQueuesTheHeldStateMessages(t *testing.T) {
	t.Parallel()

	const (
		held    = 2
		lines   = held + 1
		maxWait = 5 * time.Second
	)

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			_, _ = io.WriteString(w, stateDownLine+stateDownLine+updateLine)

			flush(w)
			<-r.Context().Done()
		}))
	defer srv.Close()

	s := newStreamerWithFullStateQueue(srv.URL)
	// The raw handler sees each line before the read loop hands it to the
	// queues, so once it has seen the UPDATE, both STATE messages are held.
	linesSeen := make(chan struct{}, lines)

	s.RegisterRawHandler(func(string) { linesSeen <- struct{}{} })

	// As in TestStopBeforeMessageReachesQueues, setting cancel as Start does
	// lets Stop cancel the stream run here.
	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel
	streamErr := make(chan error, 1)

	go func() { streamErr <- s.stream(ctx) }()

	for range lines {
		select {
		case <-linesSeen:
		case <-time.After(maxWait):
			t.Fatalf("the read loop did not read %d lines within %s",
				lines, maxWait)
		}
	}

	// No worker runs here: the test takes from the full queue in its place.
	queue := s.handlers[0].queue
	taken := takeAll(queue)
	stopped := make(chan struct{})

	go func() {
		s.Stop(t.Context())
		close(stopped)
	}()

	select {
	case <-stopped:
	case <-time.After(maxWait):
		t.Fatalf("Stop did not return within %s while a STATE message was held",
			maxWait)
	}

	// Stop has closed the queue, so takeAll is done or about to be.
	msgs := <-taken
	if len(msgs) != cap(queue)+held {
		t.Fatalf("%d messages came off the queue, want the %d on it and the "+
			"%d held", len(msgs), cap(queue), held)
	}

	for _, msg := range msgs[cap(queue):] {
		if msg.Peer != stateDownPeer || msg.State != ristypes.StateDown {
			t.Errorf("message after those on the queue has peer %q, state %q; "+
				"want the down from %s", msg.Peer, msg.State, stateDownPeer)
		}
	}

	select {
	case err := <-streamErr:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("stream returned %v, want %v", err, context.Canceled)
		}
	case <-time.After(maxWait):
		t.Fatalf("stream did not end within %s of Stop", maxWait)
	}
}

// TestStopReturnsByItsDeadline stops a streamer whose handler has a full queue
// and a STATE message held for it, and a worker that stays busy until the test
// ends, as one does whose handler waits for a database write. Stop must return
// within its deadline plus a small margin instead of waiting for room on the
// queue and for the worker, and log as not handed over the held message and
// every message on the queue.
func TestStopReturnsByItsDeadline(t *testing.T) {
	t.Parallel()

	const (
		deadline = 100 * time.Millisecond
		margin   = 2 * time.Second
	)

	var logs bytes.Buffer

	s := newStreamerWithFullStateQueue("")
	s.logger = &logger.Logger{Logger: slog.New(slog.NewJSONHandler(&logs,
		&slog.HandlerOptions{Level: slog.LevelError}))}
	s.handlers[0].heldState = []*ristypes.RISMessage{
		{Type: ristypes.MessageTypeState},
	}
	wantNotHandedOver := 1 + len(s.handlers[0].queue)
	// Stop cancels the stream Start would have started; there is none here.
	s.cancel = func() {}

	release := make(chan struct{})
	defer close(release)

	s.workers.Go(func() { <-release })

	ctx, cancel := context.WithTimeout(t.Context(), deadline)
	defer cancel()

	stopped := make(chan struct{})

	go func() {
		s.Stop(ctx)
		close(stopped)
	}()

	select {
	case <-stopped:
	case <-time.After(deadline + margin):
		t.Fatalf("Stop did not return within %s of its deadline", margin)
	}

	var logged struct {
		Messages int `json:"messages"`
	}

	err := json.Unmarshal(logs.Bytes(), &logged)
	if err != nil {
		t.Fatalf("want the error Stop logs, got %q: %v", logs.String(), err)
	}

	if logged.Messages != wantNotHandedOver {
		t.Errorf("Stop logged %d messages not handed over, want %d",
			logged.Messages, wantNotHandedOver)
	}
}

// TestStopBeforeMessageReachesQueues stops the streamer after the read loop
// has checked for cancellation but before it hands the message to the handler
// queues. That is the gap a stop from another goroutine can land in, and it
// used to end in "send on closed channel". The raw handler runs in that gap on
// the read loop itself, so calling Stop from it hits the gap every time.
func TestStopBeforeMessageReachesQueues(t *testing.T) {
	t.Parallel()

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, updateLine)
		}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL
	s.RegisterHandler(updateHandler{})
	s.RegisterRawHandler(func(string) { s.Stop(t.Context()) })

	// Start would run the stream in the background, where the test cannot
	// wait for it. Setting cancel as Start does lets Stop cancel the stream
	// run here instead.
	ctx, cancel := context.WithCancel(context.Background())
	s.cancel = cancel

	err := s.stream(ctx)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("stream returned %v, want %v", err, context.Canceled)
	}

	// A second Stop must not close the queues again.
	s.Stop(t.Context())
}

// TestReconnectWhenStreamEnds has the server end the first response
// normally, as when RIS Live restarts. The streamer must treat that as a
// disconnect and connect again.
func TestReconnectWhenStreamEnds(t *testing.T) {
	t.Parallel()

	endResponse := func(http.ResponseWriter, *http.Request) {}
	checkReconnect(t, endResponse, streamIdleTimeout)
}

// TestReconnectWhenStreamGoesIdle has the server hold the first response
// open without sending anything. The streamer must drop it once the idle
// timeout passes and connect again.
func TestReconnectWhenStreamGoesIdle(t *testing.T) {
	t.Parallel()

	const idleTimeout = 50 * time.Millisecond

	holdResponse := func(w http.ResponseWriter, r *http.Request) {
		flush(w)
		<-r.Context().Done()
	}
	checkReconnect(t, holdResponse, idleTimeout)
}

// TestReconnectWhenServerNeverAnswers has the server accept the first
// request and never answer it, not even with headers. The streamer must give
// up on it once the idle timeout passes and connect again.
func TestReconnectWhenServerNeverAnswers(t *testing.T) {
	t.Parallel()

	const idleTimeout = 50 * time.Millisecond

	neverAnswer := func(_ http.ResponseWriter, r *http.Request) {
		<-r.Context().Done()
	}
	checkReconnect(t, neverAnswer, idleTimeout)
}

// TestStreamNotDroppedWhileLinesArrive has the server send lines well within
// the idle timeout for several times the timeout, then end the response. Each
// line must start the idle timeout over, so the stream lasts until the server
// ends it instead of being dropped as idle.
func TestStreamNotDroppedWhileLinesArrive(t *testing.T) {
	t.Parallel()

	const (
		idleTimeout  = 100 * time.Millisecond
		lineInterval = idleTimeout / 10
		sendFor      = 5 * idleTimeout
	)

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			deadline := time.Now().Add(sendFor)
			for time.Now().Before(deadline) {
				_, _ = io.WriteString(w, updateLine)

				flush(w)
				time.Sleep(lineInterval)
			}
		}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL
	s.idleTimeout = idleTimeout

	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
	}
}

// TestStreamReadsLongLine has the server send one UPDATE line of about
// 200 KiB, padded with raw hex the way RIS Live writes an extended BGP
// message, and then end the response. The line is longer than a
// bufio.Scanner reads by default. The streamer must hand the message to the
// handler's queue and read on until the server ends the stream, instead of
// dropping the connection at the long line.
func TestStreamReadsLongLine(t *testing.T) {
	t.Parallel()

	const rawHexLength = 200 * bytesPerKB

	longLine := `{"type":"ris_message","data":{"type":"UPDATE","peer":"192.0.2.1",` +
		`"peer_asn":"64496","timestamp":1700000000,"raw":"` +
		strings.Repeat("0", rawHexLength) + `"}}` + "\n"

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, longLine)
		}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL
	s.RegisterHandler(updateHandler{})

	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
	}

	queued := len(s.handlers[0].queue)
	if queued != 1 {
		t.Errorf("handler queue holds %d messages, want 1", queued)
	}
}

// blockedHandler wants UPDATE messages, has room on its queue for one, and
// does not return from HandleMessage until release is closed.
type blockedHandler struct {
	release chan struct{}
}

func (blockedHandler) WantsMessage(messageType string) bool {
	return messageType == ristypes.MessageTypeUpdate
}

func (blockedHandler) QueueCapacity() int { return 1 }

func (h blockedHandler) HandleMessage(*ristypes.RISMessage) { <-h.release }

// TestHandlerStatsReadWhileMessagesAreDropped has the server send a few UPDATE
// lines to a blockedHandler, so the read loop drops all but the one the
// handler holds and the one on its queue, while the test reads the handler's
// statistics. The drop count must reach the number dropped, and the race
// detector must find the count read and written safely.
func TestHandlerStatsReadWhileMessagesAreDropped(t *testing.T) {
	t.Parallel()

	const (
		lines       = 5
		wantDropped = lines - 2
		maxWait     = 5 * time.Second
	)

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			_, _ = io.WriteString(w, strings.Repeat(updateLine, lines))

			flush(w)
			<-r.Context().Done()
		}))
	defer srv.Close()

	s := New(logger.New(), metrics.New())
	s.url = srv.URL
	handler := blockedHandler{release: make(chan struct{})}
	s.RegisterHandler(handler)

	err := s.Start(t.Context())
	if err != nil {
		t.Fatalf("Start returned error: %v", err)
	}
	defer s.Stop(t.Context())
	defer close(handler.release)

	// The test polls instead of waiting for a signal from the read loop: a
	// signal would order the drops before the reads here, and the race
	// detector would have nothing to report.
	deadline := time.Now().Add(maxWait)

	for s.GetHandlerStats()[0].DroppedCount < wantDropped {
		if time.Now().After(deadline) {
			t.Fatalf("fewer than %d messages counted as dropped within %s",
				wantDropped, maxWait)
		}
	}
}

// TestPeerCountCountsEveryPeerHeardFrom has the server send an UPDATE from
// stateDownPeer and a KEEPALIVE from another peer, with no OPEN from either,
// and then a message that ends stateDownPeer's session: a NOTIFICATION, or a
// STATE message whose state is down. The peer count must be 2 after the first
// two messages and 1 after the third.
func TestPeerCountCountsEveryPeerHeardFrom(t *testing.T) {
	t.Parallel()

	const (
		otherPeer      = "192.0.2.2"
		peerUpdateLine = `{"type":"ris_message","data":{"type":"UPDATE",` +
			`"peer":"` + stateDownPeer + `",` +
			`"peer_asn":"64496","timestamp":1700000000}}` + "\n"
		otherPeerKeepaliveLine = `{"type":"ris_message","data":{"type":"KEEPALIVE",` +
			`"peer":"` + otherPeer + `",` +
			`"peer_asn":"64497","timestamp":1700000000}}` + "\n"
		peerNotificationLine = `{"type":"ris_message","data":{"type":"NOTIFICATION",` +
			`"peer":"` + stateDownPeer + `",` +
			`"peer_asn":"64496","timestamp":1700000000}}` + "\n"
		bothPeers = 2
	)

	sessionEndLines := map[string]string{
		"NOTIFICATION": peerNotificationLine,
		"STATE down":   stateDownLine,
	}

	for name, sessionEndLine := range sessionEndLines {
		t.Run(name, func(t *testing.T) {
			t.Parallel()

			tracker := metrics.New()
			s := New(logger.New(), tracker)

			streamLines(t, s, peerUpdateLine+otherPeerKeepaliveLine)

			got := tracker.GetBGPPeerCount()
			if got != bothPeers {
				t.Fatalf("peer count is %d after messages from two peers, want %d",
					got, bothPeers)
			}

			streamLines(t, s, sessionEndLine)

			got = tracker.GetBGPPeerCount()
			if got != 1 {
				t.Errorf("peer count is %d after one of two peers' sessions "+
					"ended, want 1", got)
			}
		})
	}
}

// streamLines has a server send lines as its whole response and runs one
// stream of s against it.
func streamLines(t *testing.T, s *Streamer, lines string) {
	t.Helper()

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, _ *http.Request) {
			_, _ = io.WriteString(w, lines)
		}))
	defer srv.Close()

	s.url = srv.URL

	err := s.stream(context.Background())
	if !errors.Is(err, errStreamEnded) {
		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
	}
}

// checkReconnect runs a streamer against a fake RIS Live server. The first
// response is served by firstResponse; every later one carries one UPDATE
// line and then stays open. The streamer must come back for a second
// response, must not be marked connected when it does, and must handle the
// UPDATE in it.
func checkReconnect(
	t *testing.T, firstResponse http.HandlerFunc, idleTimeout time.Duration,
) {
	t.Helper()

	const maxWait = 5 * time.Second

	tracker := metrics.New()
	connectedAtReconnect := make(chan bool, 1)

	var requests atomic.Int32

	srv := httptest.NewServer(http.HandlerFunc(
		func(w http.ResponseWriter, r *http.Request) {
			if requests.Add(1) == 1 {
				firstResponse(w, r)

				return
			}

			select {
			case connectedAtReconnect <- tracker.IsConnected():
			default:
			}

			_, _ = io.WriteString(w, updateLine)

			flush(w)
			<-r.Context().Done()
		}))
	defer srv.Close()

	s := New(logger.New(), tracker)
	s.url = srv.URL
	s.idleTimeout = idleTimeout
	s.minBackoff = time.Millisecond
	handler := updateHandler{handled: make(chan struct{}, 1)}
	s.RegisterHandler(handler)

	err := s.Start(t.Context())
	if err != nil {
		t.Fatalf("Start returned error: %v", err)
	}
	defer s.Stop(t.Context())

	select {
	case <-handler.handled:
	case <-time.After(maxWait):
		t.Fatalf("no message from a second response handled within %s", maxWait)
	}

	if <-connectedAtReconnect {
		t.Error("streamer was still marked connected when it came back " +
			"for a second response")
	}
}

// settledGoroutineCount lets transient goroutines finish, then reports the
// current count.
func settledGoroutineCount() int {
	prev := runtime.NumGoroutine()

	for range 20 {
		time.Sleep(10 * time.Millisecond)

		cur := runtime.NumGoroutine()
		if cur == prev {
			return cur
		}

		prev = cur
	}

	return prev
}

// waitForGoroutines waits until the goroutine count drops to target or below.
func waitForGoroutines(target int) bool {
	for range 100 {
		if runtime.NumGoroutine() <= target {
			return true
		}

		time.Sleep(10 * time.Millisecond)
	}

	return false
}