routewatch/internal/database/database_internal_test.go

package database

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"net/netip"
	"slices"
	"sync"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/logger"
	"github.com/google/uuid"
	"github.com/mattn/go-sqlite3"
)

// tempStoreMemory is the PRAGMA temp_store value meaning "hold temp B-trees in
// memory"; the DSN change must leave temp_store below this so they spill to disk.
const tempStoreMemory = 2

// autoVacuumIncremental is the PRAGMA auto_vacuum value meaning "incremental".
const autoVacuumIncremental = 2

// vacuumTestRoutes is how many routes the vacuum test writes and then deletes,
// enough to leave many free pages in the file.
const vacuumTestRoutes = 2000

// heldConnections is how many pooled connections the pragma test holds open at
// once so each is a distinct SQLite connection that parsed the DSN.
const heldConnections = 5

// testPeerIP is the peer every route in the IP lookup test is learned from.
const testPeerIP = "192.0.2.254"

// Parameters for the checkpoint-contention regression test.
const (
	// contentionIterations is how many batch writes race the checkpoint loop.
	contentionIterations = 400
	// contendedASNCount is the small set of ASNs the batches reuse, so most
	// batches update existing rows and exercise the read-before-write path.
	contendedASNCount = 16
	// asnSecondBand offsets a second ASN per batch so each batch writes more
	// than one row.
	asnSecondBand = 100
)

// Parameters for the prefix-length sample tests.
const (
	// sampleLimit is how many prefixes the tests ask the sample for.
	sampleLimit = 10
	// sampledPrefixCount is how many prefixes the sampling test stores, more
	// than sampleLimit so the sample has to choose among them.
	sampledPrefixCount = 30
	// secondV4PeerIP and secondV6PeerIP are the second peer each prefix is
	// learned from, besides the one mkV4Route and mkV6Route set.
	secondV4PeerIP = "192.0.2.2"
	secondV6PeerIP = "2001:db8::2"
	// largeSamplePrefixCount is how many prefixes the test of a mask length
	// with more than maxPrefixesToShuffle prefixes stores: twice that, so about
	// half of them come after the first maxPrefixesToShuffle+1 in text order.
	largeSamplePrefixCount = 2 * maxPrefixesToShuffle
	// largeSampleLimit is how many prefixes that test asks for. Each random
	// lookup lands after the first maxPrefixesToShuffle+1 in text order about
	// half the time, so none of them doing so is about as likely as
	// largeSampleLimit coin tosses all coming up heads.
	largeSampleLimit = 30
	// liveRoutesV4Table and liveRoutesV6Table are the live route tables the
	// tests pass to the sampling steps.
	liveRoutesV4Table = "live_routes_v4"
	liveRoutesV6Table = "live_routes_v6"
)

// TestGetIPInfoFindsMostSpecificLiveRoute stores nested live prefixes for both
// families and checks that a lookup returns the most specific one covering the
// address, ErrNoRoute when none covers it, and the next less specific prefix
// once the only route of the most specific one is withdrawn.
func TestGetIPInfoFindsMostSpecificLiveRoute(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	// Nested live prefixes, each originated by its own AS.
	origins := map[string]int{
		"10.0.0.0/8":        64500,
		"10.1.0.0/16":       64501,
		"10.1.2.0/24":       64502,
		"2001:db8::/32":     64500,
		"2001:db8:1::/48":   64501,
		"2001:db8:1:2::/64": 64502,
	}
	writeOriginRoutes(t, db, origins)

	// lookup checks that ip resolves to the live prefix want, or to ErrNoRoute
	// when want is empty.
	lookup := func(ip, want string) {
		t.Helper()

		info, err := db.GetIPInfo(ip)
		if want == "" {
			if !errors.Is(err, ErrNoRoute) {
				t.Errorf("GetIPInfo(%s) = %+v, %v; want ErrNoRoute", ip, info, err)
			}

			return
		}

		if err != nil {
			t.Errorf("GetIPInfo(%s): %v", ip, err)

			return
		}

		if info.Netblock != want || info.MaskLength != netip.MustParsePrefix(want).Bits() ||
			info.ASN != origins[want] {
			t.Errorf("GetIPInfo(%s) = %s (mask %d) AS%d, want %s AS%d",
				ip, info.Netblock, info.MaskLength, info.ASN, want, origins[want])
		}
	}

	lookup("10.1.2.3", "10.1.2.0/24")
	lookup("::ffff:10.1.2.3", "10.1.2.0/24")
	lookup("10.1.3.4", "10.1.0.0/16")
	lookup("10.2.0.1", "10.0.0.0/8")
	lookup("192.0.2.1", "")
	lookup("2001:db8:1:2::3", "2001:db8:1:2::/64")
	lookup("2001:db8:1:3::4", "2001:db8:1::/48")
	lookup("2001:db8:2::1", "2001:db8::/32")
	lookup("2001:db9::1", "")

	err = db.DeleteLiveRouteBatch([]LiveRouteDeletion{
		{Prefix: "10.1.2.0/24", PeerIP: testPeerIP, IPVersion: ipVersionV4},
		{Prefix: "2001:db8:1:2::/64", PeerIP: testPeerIP, IPVersion: ipVersionV6},
	})
	if err != nil {
		t.Fatalf("DeleteLiveRouteBatch: %v", err)
	}

	lookup("10.1.2.3", "10.1.0.0/16")
	lookup("2001:db8:1:2::3", "2001:db8:1::/48")
}

// TestGetIPInfoIgnoresShortRoutes checks that the IP lookup does not report a
// default route, or a route shorter than /8 for IPv4 or /16 for IPv6, as the
// route for an address, and that a route of exactly /8 or /16, or longer, still
// is. The prefixes are written out so that a wrong limit fails the test.
func TestGetIPInfoIgnoresShortRoutes(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	// announce writes one live route for prefix, originated by asn.
	announce := func(prefix string, asn int) {
		t.Helper()

		writeOriginRoutes(t, db, map[string]int{prefix: asn})
	}

	// lookup checks that ip resolves to the live prefix netblock from asn.
	lookup := func(ip, netblock string, asn int) {
		t.Helper()

		info, err := db.GetIPInfo(ip)
		if err != nil {
			t.Errorf("GetIPInfo(%s): %v", ip, err)

			return
		}

		if info.Netblock != netblock || info.ASN != asn {
			t.Errorf("GetIPInfo(%s) = %s AS%d, want %s AS%d",
				ip, info.Netblock, info.ASN, netblock, asn)
		}
	}

	// Default and very short routes, and routes one bit shorter than each limit.
	announce("2000::/3", 64500)
	announce("0.0.0.0/0", 64500)
	announce("2000::/15", 64500)
	announce("192.0.0.0/7", 64500)

	for _, ip := range []string{"2001:db8::1", "192.0.2.1"} {
		info, err := db.GetIPInfo(ip)
		if !errors.Is(err, ErrNoRoute) {
			t.Errorf("GetIPInfo(%s) with only short routes = %+v, %v; "+
				"want ErrNoRoute", ip, info, err)
		}
	}

	// Routes of exactly the minimum length are used.
	announce("2001::/16", 64501)
	announce("192.0.0.0/8", 64501)

	lookup("2001:db8::1", "2001::/16", 64501)
	lookup("192.0.2.1", "192.0.0.0/8", 64501)

	announce("2001:db8::/32", 64502)

	lookup("2001:db8::1", "2001:db8::/32", 64502)
}

// TestGetIPInfoNeedsNoWHOISRefreshForOriginZero checks that the IP lookup asks
// for no WHOIS lookup when the address's route has origin 0, no known origin,
// and still asks for one when its origin AS has no WHOIS data.
func TestGetIPInfoNeedsNoWHOISRefreshForOriginZero(t *testing.T) {
	t.Parallel()

	db, err := New(&config.Config{StateDir: t.TempDir()}, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	writeOriginRoutes(t, db, map[string]int{"10.0.0.0/8": 0, "11.0.0.0/8": 64500})

	for ip, want := range map[string]bool{"10.0.0.1": false, "11.0.0.1": true} {
		info, err := db.GetIPInfo(ip)
		if err != nil {
			t.Fatalf("GetIPInfo(%s): %v", ip, err)
		}

		if info.NeedsWHOISRefresh != want {
			t.Errorf("GetIPInfo(%s) asks for a WHOIS lookup of AS%d: %t, want %t",
				ip, info.ASN, info.NeedsWHOISRefresh, want)
		}
	}
}

// TestLiveRouteTablesHoldOneRoutePerPrefixAndPeer checks that each live route
// table refuses a second route from a peer for a prefix, even with another
// origin AS. An upsert updates the peer's route before it would insert one, so
// only an insert straight into the table shows what the table allows.
func TestLiveRouteTablesHoldOneRoutePerPrefixAndPeer(t *testing.T) {
	t.Parallel()

	db, err := New(&config.Config{StateDir: t.TempDir()}, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	const otherOriginASN = 64501

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)

	for _, tt := range []struct {
		route     *LiveRoute
		insertSQL string
	}{
		{
			route:     mkV4Route(t, "198.51.100.0/24", 64500, ts),
			insertSQL: insertLiveRouteV4SQL,
		},
		{
			route:     mkV6Route(testV6Prefix, 64500, ts),
			insertSQL: insertLiveRouteV6SQL,
		},
	} {
		err := db.UpsertLiveRoute(tt.route)
		if err != nil {
			t.Fatalf("UpsertLiveRoute(%s): %v", tt.route.Prefix, err)
		}

		_, err = db.db.ExecContext(t.Context(), tt.insertSQL, uuid.NewString(),
			tt.route.Prefix, tt.route.MaskLength, otherOriginASN,
			tt.route.PeerIP, "[]", tt.route.NextHop, ts)

		var sqliteErr sqlite3.Error
		if !errors.As(err, &sqliteErr) ||
			sqliteErr.ExtendedCode != sqlite3.ErrConstraintUnique {
			t.Errorf("second route from peer %s for %s: error %v, "+
				"want a UNIQUE constraint error",
				tt.route.PeerIP, tt.route.Prefix, err)
		}
	}
}

// writeOriginRoutes writes one live route from testPeerIP for each prefix in
// origins, originated by the AS it maps to.
func writeOriginRoutes(t *testing.T, db *Database, origins map[string]int) {
	t.Helper()

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	routes := make([]*LiveRoute, 0, len(origins))

	for prefix, asn := range origins {
		routes = append(routes, &LiveRoute{
			ID:          uuid.New(),
			Prefix:      prefix,
			MaskLength:  netip.MustParsePrefix(prefix).Bits(),
			IPVersion:   detectIPVersion(prefix),
			OriginASN:   asn,
			PeerIP:      testPeerIP,
			ASPath:      []int{asn},
			NextHop:     testPeerIP,
			LastUpdated: ts,
		})
	}

	err := db.UpsertLiveRouteBatch(routes)
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}
}

// TestConnectionPoolPragmas holds several pooled connections open at once and
// checks each one carries the per-connection settings from the DSN, plus the
// process-wide hard heap limit, and sees the new file with auto_vacuum
// incremental.
func TestConnectionPoolPragmas(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()

	// Hold distinct connections open simultaneously so the pool must open a new
	// one (each parsing the DSN) rather than hand back the same connection.
	conns := make([]*sql.Conn, 0, heldConnections)
	defer func() {
		for _, c := range conns {
			_ = c.Close()
		}
	}()

	for i := range heldConnections {
		c, err := db.db.Conn(ctx)
		if err != nil {
			t.Fatalf("failed to open connection %d: %v", i, err)
		}

		conns = append(conns, c)
	}

	for i, c := range conns {
		cacheSize := readPragma(ctx, t, c, i, "cache_size")
		if cacheSize != sqliteCacheSizeKiB {
			t.Errorf("conn %d: cache_size = %d, want %d",
				i, cacheSize, sqliteCacheSizeKiB)
		}

		busyTimeout := readPragma(ctx, t, c, i, "busy_timeout")
		if busyTimeout != sqliteBusyTimeoutMs {
			t.Errorf("conn %d: busy_timeout = %d, want %d",
				i, busyTimeout, sqliteBusyTimeoutMs)
		}

		tempStore := readPragma(ctx, t, c, i, "temp_store")
		if tempStore == tempStoreMemory {
			t.Errorf("conn %d: temp_store = %d, want anything but %d (MEMORY)",
				i, tempStore, tempStoreMemory)
		}

		hardHeapLimit := readPragma(ctx, t, c, i, "hard_heap_limit")
		if hardHeapLimit != sqliteHardHeapLimitBytes {
			t.Errorf("conn %d: hard_heap_limit = %d, want %d",
				i, hardHeapLimit, sqliteHardHeapLimitBytes)
		}

		autoVacuum := readPragma(ctx, t, c, i, "auto_vacuum")
		if autoVacuum != autoVacuumIncremental {
			t.Errorf("conn %d: auto_vacuum = %d, want %d (incremental)",
				i, autoVacuum, autoVacuumIncremental)
		}
	}
}

// readPragma reads the integer value of PRAGMA name on connection i, c.
func readPragma(
	ctx context.Context, t *testing.T, c *sql.Conn, i int, name string,
) int64 {
	t.Helper()

	var value int64

	err := c.QueryRowContext(ctx, "PRAGMA "+name).Scan(&value)
	if err != nil {
		t.Fatalf("conn %d: failed to read %s: %v", i, name, err)
	}

	return value
}

// TestVacuumReturnsFreePages checks that after routes are deleted from a new
// database, one Vacuum call returns every page they used. The deletes leave
// fewer free pages than the 1000 Vacuum frees per call, so none may remain.
// With auto_vacuum off (issue https://git.eeqj.de/sneak/routewatch/issues/43)
// the free pages stayed in the file, and with the PRAGMA run by ExecContext
// Vacuum freed only one page per call.
func TestVacuumReturnsFreePages(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ctx := context.Background()
	ts := time.Now().UTC()

	const asn = 64500

	routes := make([]*LiveRoute, 0, vacuumTestRoutes)
	deletions := make([]LiveRouteDeletion, 0, vacuumTestRoutes)

	for i := range vacuumTestRoutes {
		route := mkV4Route(t, fmt.Sprintf("10.%d.%d.0/24", i/256, i%256), asn, ts)
		routes = append(routes, route)
		deletions = append(deletions, LiveRouteDeletion{
			Prefix:    route.Prefix,
			PeerIP:    route.PeerIP,
			IPVersion: ipVersionV4,
		})
	}

	err = db.UpsertLiveRouteBatch(routes)
	if err != nil {
		t.Fatalf("failed to write routes: %v", err)
	}

	err = db.DeleteLiveRouteBatch(deletions)
	if err != nil {
		t.Fatalf("failed to delete routes: %v", err)
	}

	var before int

	err = db.db.QueryRowContext(ctx, "PRAGMA freelist_count").Scan(&before)
	if err != nil {
		t.Fatalf("failed to read freelist_count: %v", err)
	}

	if before == 0 {
		t.Fatalf("no free pages after deleting %d routes", vacuumTestRoutes)
	}

	err = db.Vacuum(ctx)
	if err != nil {
		t.Fatalf("Vacuum failed: %v", err)
	}

	var after int

	err = db.db.QueryRowContext(ctx, "PRAGMA freelist_count").Scan(&after)
	if err != nil {
		t.Fatalf("failed to read freelist_count: %v", err)
	}

	if after != 0 {
		t.Errorf("free pages after Vacuum = %d of %d, want 0", after, before)
	}
}

// TestBatchWriteDuringCheckpoint reproduces issue #25. A batch write reads
// (SELECT) before it writes (INSERT/UPDATE). Under the default deferred locking
// the transaction begins as a reader and, when the maintainer's WAL checkpoint
// holds the write lock, its upgrade to writer fails immediately with "database
// is locked" without honouring busy_timeout, dropping the batch. With
// _txlock=immediate the transaction takes the write lock at BEGIN and waits, so
// no batch is dropped. The checkpoint runs without the Database mutex, exactly
// as the background maintainer does in production.
func TestBatchWriteDuringCheckpoint(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ctx, cancel := context.WithCancel(context.Background())

	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()

		for {
			select {
			case <-ctx.Done():
				return
			default:
				_ = db.Checkpoint(ctx) // errors are the checkpoint's own to absorb
			}
		}
	}()

	ts := time.Now().UTC()
	for i := range contentionIterations {
		asns := map[int]time.Time{
			i % contendedASNCount:                   ts,
			(i % contendedASNCount) + asnSecondBand: ts,
		}

		err := db.GetOrCreateASNBatch(asns)
		if err != nil {
			cancel()
			wg.Wait()
			t.Fatalf("batch write failed under checkpoint contention: %v", err)
		}
	}

	cancel()
	wg.Wait()
}

// TestGetRandomPrefixesByLengthListsPrefixOnce stores one prefix of each
// family learned from two peers in the same second and checks that the
// prefix-length sample lists each of them once.
func TestGetRandomPrefixesByLengthListsPrefixOnce(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	v4First := mkV4Route(t, "198.51.100.0/24", 64500, ts)
	v4Second := mkV4Route(t, "198.51.100.0/24", 64500, ts)
	v4Second.PeerIP = secondV4PeerIP
	v6First := mkV6Route(testV6Prefix, 64500, ts)
	v6Second := mkV6Route(testV6Prefix, 64500, ts)
	v6Second.PeerIP = secondV6PeerIP

	err = db.UpsertLiveRouteBatch(
		[]*LiveRoute{v4First, v4Second, v6First, v6Second})
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	for _, want := range []*LiveRoute{v4First, v6First} {
		routes, err := db.GetRandomPrefixesByLength(want.MaskLength,
			want.IPVersion, sampleLimit)
		if err != nil {
			t.Fatalf("GetRandomPrefixesByLength(%d, %d): %v",
				want.MaskLength, want.IPVersion, err)
		}

		if len(routes) != 1 || routes[0].Prefix != want.Prefix {
			t.Errorf("GetRandomPrefixesByLength(%d, %d) = %+v, want %s listed once",
				want.MaskLength, want.IPVersion, routes, want.Prefix)
		}
	}
}

// TestGetRandomPrefixesByLengthListsNewestRoute stores a prefix learned from
// two peers at different times, the older route first, and checks that the
// prefix-length sample lists the newer route.
func TestGetRandomPrefixesByLengthListsNewestRoute(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	older := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	olderRoute := mkV4Route(t, "198.51.100.0/24", 64500, older)
	newerRoute := mkV4Route(t, "198.51.100.0/24", 64500, older.Add(time.Minute))
	newerRoute.PeerIP = secondV4PeerIP

	err = db.UpsertLiveRouteBatch([]*LiveRoute{olderRoute, newerRoute})
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	routes, err := db.GetRandomPrefixesByLength(newerRoute.MaskLength,
		ipVersionV4, sampleLimit)
	if err != nil {
		t.Fatalf("GetRandomPrefixesByLength: %v", err)
	}

	if len(routes) != 1 || routes[0].PeerIP != newerRoute.PeerIP {
		t.Errorf("GetRandomPrefixesByLength = %+v, want only the newer route, "+
			"from %s", routes, newerRoute.PeerIP)
	}
}

// TestGetRandomPrefixesByLengthSamplesManyPrefixes stores more /24 prefixes
// than the sample asks for, each followed by a /16 so that routes of another
// mask length sit between the /24 rows, and each prefix learned from two peers
// in the same second. It checks both the sample the pages get, picked from the
// shuffled prefixes, and the random lookups used for a mask length with more
// than maxPrefixesToShuffle prefixes.
func TestGetRandomPrefixesByLengthSamplesManyPrefixes(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)

	var routes []*LiveRoute

	stored := make([]string, 0, sampledPrefixCount)

	for i := range sampledPrefixCount {
		prefix := fmt.Sprintf("10.0.%d.0/24", i)
		for _, p := range []string{prefix, fmt.Sprintf("172.%d.0.0/16", i)} {
			first := mkV4Route(t, p, 64500, ts)
			second := mkV4Route(t, p, 64500, ts)
			second.PeerIP = secondV4PeerIP
			routes = append(routes, first, second)
		}

		stored = append(stored, prefix)
	}

	err = db.UpsertLiveRouteBatch(routes)
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	slices.Sort(stored)

	maskLength := routes[0].MaskLength

	sample, err := db.GetRandomPrefixesByLength(maskLength, ipVersionV4, sampleLimit)
	if err != nil {
		t.Fatalf("GetRandomPrefixesByLength: %v", err)
	}

	listed := make([]string, 0, len(sample))
	for _, route := range sample {
		listed = append(listed, route.Prefix)
	}

	checkPrefixSample(t, "GetRandomPrefixesByLength", listed, stored)

	looked, err := db.samplePrefixesByLength(
		context.Background(), liveRoutesV4Table, maskLength, sampleLimit, nil)
	if err != nil {
		t.Fatalf("samplePrefixesByLength: %v", err)
	}

	checkPrefixSample(t, "samplePrefixesByLength", looked, stored)
}

// checkPrefixSample checks that sample lists sampleLimit prefixes, each once
// and each one of stored, and that it is not just the first sampleLimit of
// stored, which is in text order.
func checkPrefixSample(t *testing.T, name string, sample, stored []string) {
	t.Helper()

	if len(sample) != sampleLimit {
		t.Errorf("%s listed %d prefixes, want %d", name, len(sample), sampleLimit)
	}

	listed := make(map[string]bool)

	for _, prefix := range sample {
		if listed[prefix] || !slices.Contains(stored, prefix) {
			t.Errorf("%s lists %s twice or it is not one of the /24s", name, prefix)
		}

		listed[prefix] = true
	}

	if slices.Equal(slices.Sorted(slices.Values(sample)), stored[:sampleLimit]) {
		t.Errorf("%s lists just the first %d prefixes in text order",
			name, sampleLimit)
	}
}

// TestGetRandomPrefixesByLengthSamplesPastFirstPrefixes stores more prefixes
// of one mask length than maxPrefixesToShuffle. It checks that
// firstPrefixesByLength reads just the first maxPrefixesToShuffle+1 in text
// order, and that the sample lists the prefixes asked for, each once, some of
// them after those first ones, which only the random lookups reach.
func TestGetRandomPrefixesByLengthSamplesPastFirstPrefixes(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	routes := make([]*LiveRoute, 0, largeSamplePrefixCount)
	stored := make([]string, 0, largeSamplePrefixCount)

	for i := range largeSamplePrefixCount {
		route := mkV6Route(fmt.Sprintf("2001:%x::/32", i), 64500, ts)
		routes = append(routes, route)
		stored = append(stored, route.Prefix)
	}

	err = db.UpsertLiveRouteBatch(routes)
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	slices.Sort(stored)
	firstPrefixes := stored[:maxPrefixesToShuffle+1]
	maskLength := routes[0].MaskLength

	read, err := db.firstPrefixesByLength(context.Background(),
		liveRoutesV6Table, maskLength, maxPrefixesToShuffle+1)
	if err != nil {
		t.Fatalf("firstPrefixesByLength: %v", err)
	}

	if !slices.Equal(read, firstPrefixes) {
		t.Errorf("firstPrefixesByLength read %d prefixes, want the first %d "+
			"in text order", len(read), len(firstPrefixes))
	}

	sample, err := db.GetRandomPrefixesByLength(maskLength, ipVersionV6,
		largeSampleLimit)
	if err != nil {
		t.Fatalf("GetRandomPrefixesByLength: %v", err)
	}

	if len(sample) != largeSampleLimit {
		t.Errorf("GetRandomPrefixesByLength listed %d prefixes, want %d",
			len(sample), largeSampleLimit)
	}

	listed := make(map[string]bool)
	pastFirstPrefixes := 0

	for _, route := range sample {
		if listed[route.Prefix] {
			t.Errorf("GetRandomPrefixesByLength lists %s twice", route.Prefix)
		}

		listed[route.Prefix] = true

		if !slices.Contains(firstPrefixes, route.Prefix) {
			pastFirstPrefixes++
		}
	}

	if pastFirstPrefixes == 0 {
		t.Errorf("GetRandomPrefixesByLength lists only prefixes among the first "+
			"%d in text order", len(firstPrefixes))
	}
}

// TestSamplePrefixesByLengthFillsUpFromShuffled stores a single prefix, so
// every random lookup lands on it, and checks that the sample still lists
// sampleLimit prefixes: the stored one, then the shuffled ones in order. The
// stored prefix is the second shuffled one, so the fill-up must skip it.
func TestSamplePrefixesByLengthFillsUpFromShuffled(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	shuffled := make([]string, 0, sampledPrefixCount)
	for i := range sampledPrefixCount {
		shuffled = append(shuffled, fmt.Sprintf("10.0.%d.0/24", i))
	}

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	route := mkV4Route(t, shuffled[1], 64500, ts)

	err = db.UpsertLiveRouteBatch([]*LiveRoute{route})
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	sample, err := db.samplePrefixesByLength(context.Background(),
		liveRoutesV4Table, route.MaskLength, sampleLimit, shuffled)
	if err != nil {
		t.Fatalf("samplePrefixesByLength: %v", err)
	}

	want := append([]string{shuffled[1], shuffled[0]}, shuffled[2:sampleLimit]...)
	if !slices.Equal(sample, want) {
		t.Errorf("samplePrefixesByLength = %v, want %v", sample, want)
	}
}

// TestGetRandomPrefixesByLengthSkipsWithdrawnRoutes withdraws a route after
// storing it, as can happen between the steps of GetRandomPrefixesByLength. It
// checks that the random lookups pick nothing for a mask length with no route
// left, and that only the prefixes that still have a route get their newest
// route listed.
func TestGetRandomPrefixesByLengthSkipsWithdrawnRoutes(t *testing.T) {
	t.Parallel()

	cfg := &config.Config{StateDir: t.TempDir()}

	db, err := New(cfg, logger.New())
	if err != nil {
		t.Fatalf("failed to create database: %v", err)
	}
	defer func() { _ = db.Close() }()

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	kept := mkV4Route(t, "198.51.100.0/24", 64500, ts)
	withdrawn := mkV4Route(t, "172.16.0.0/16", 64500, ts)

	err = db.UpsertLiveRouteBatch([]*LiveRoute{kept, withdrawn})
	if err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	err = db.DeleteLiveRoute(withdrawn.Prefix, withdrawn.PeerIP)
	if err != nil {
		t.Fatalf("DeleteLiveRoute: %v", err)
	}

	ctx := context.Background()

	sample, err := db.samplePrefixesByLength(ctx, liveRoutesV4Table,
		withdrawn.MaskLength, sampleLimit, []string{withdrawn.Prefix})
	if err != nil {
		t.Fatalf("samplePrefixesByLength: %v", err)
	}

	if len(sample) != 0 {
		t.Errorf("samplePrefixesByLength = %v, want none for a mask length "+
			"with no route left", sample)
	}

	routes, err := db.newestRouteOfEachPrefix(ctx, liveRoutesV4Table,
		ipVersionV4, []string{withdrawn.Prefix, kept.Prefix})
	if err != nil {
		t.Fatalf("newestRouteOfEachPrefix: %v", err)
	}

	if len(routes) != 1 || routes[0].Prefix != kept.Prefix {
		t.Errorf("newestRouteOfEachPrefix = %+v, want only the route of %s",
			routes, kept.Prefix)
	}
}