routewatch/internal/database/counts_test.go

package database

import (
	"context"
	"slices"
	"sync"
	"testing"
	"time"

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

// mkV4Route builds an IPv4 live route with its mask length and range columns
// taken from the prefix.
func mkV4Route(t *testing.T, prefix string, asn int, ts time.Time) *LiveRoute {
	t.Helper()

	start, end, err := CalculateIPv4Range(prefix)
	if err != nil {
		t.Fatalf("CalculateIPv4Range(%s): %v", prefix, err)
	}
	maskLength, err := prefixMaskLength(prefix)
	if err != nil {
		t.Fatalf("prefixMaskLength(%s): %v", prefix, err)
	}

	return &LiveRoute{
		ID:          uuid.New(),
		Prefix:      prefix,
		MaskLength:  maskLength,
		IPVersion:   ipVersionV4,
		OriginASN:   asn,
		PeerIP:      "192.0.2.1",
		ASPath:      []int{asn},
		NextHop:     "192.0.2.254",
		LastUpdated: ts,
		V4IPStart:   &start,
		V4IPEnd:     &end,
	}
}

// mkV6Route builds an IPv6 live route.
func mkV6Route(prefix string, asn int, ts time.Time) *LiveRoute {
	return &LiveRoute{
		ID:          uuid.New(),
		Prefix:      prefix,
		MaskLength:  32,
		IPVersion:   ipVersionV6,
		OriginASN:   asn,
		PeerIP:      "2001:db8::1",
		ASPath:      []int{asn},
		NextHop:     "2001:db8::ffff",
		LastUpdated: ts,
	}
}

// TestLiveCountsTrackWritesInRealtime checks that the stats counts start at
// zero, reflect each write the moment it commits (no recompute, no timer), do
// not move when a route is merely re-announced, and drop when a route is
// deleted. These counts are what /api/v1/stats reports; before this change the
// endpoint recomputed them with a COUNT(*) over each table on every request.
func TestLiveCountsTrackWritesInRealtime(t *testing.T) {
	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()

	empty, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext on empty database: %v", err)
	}
	if empty.ASNs != 0 || empty.Prefixes != 0 || empty.Peerings != 0 ||
		empty.Peers != 0 || empty.LiveRoutes != 0 {
		t.Fatalf("empty database counts nonzero: %+v", empty)
	}

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

	if err := db.GetOrCreateASNBatch(map[int]time.Time{64500: ts, 64501: ts}); err != nil {
		t.Fatalf("GetOrCreateASNBatch: %v", err)
	}
	if err := db.UpdatePrefixesBatch(map[string]time.Time{
		"198.51.100.0/24": ts,
		"2001:db8::/32":   ts,
	}); err != nil {
		t.Fatalf("UpdatePrefixesBatch: %v", err)
	}
	if err := db.UpdatePeerBatch(map[string]PeerUpdate{
		"192.0.2.1": {PeerIP: "192.0.2.1", PeerASN: 64500, MessageType: "UPDATE", Timestamp: ts},
	}); err != nil {
		t.Fatalf("UpdatePeerBatch: %v", err)
	}
	if err := db.RecordPeering(64500, 64501, ts); err != nil {
		t.Fatalf("RecordPeering: %v", err)
	}

	routes := []*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, ts),
		mkV4Route(t, "203.0.113.0/24", 64501, ts.Add(time.Minute)),
		mkV6Route("2001:db8::/32", 64502, ts.Add(2*time.Minute)),
	}
	if err := db.UpsertLiveRouteBatch(routes); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after inserts", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 2, ipv6Routes: 1, liveRoutes: 3,
	})

	// Re-announcing the same routes is an update, not an insert: counts hold.
	if err := db.UpsertLiveRouteBatch(routes); err != nil {
		t.Fatalf("UpsertLiveRouteBatch (re-announce): %v", err)
	}
	stats, err = db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after re-announce", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 2, ipv6Routes: 1, liveRoutes: 3,
	})

	// A withdrawal removes one route.
	if err := db.DeleteLiveRouteBatch([]LiveRouteDeletion{
		{Prefix: "203.0.113.0/24", OriginASN: 64501, PeerIP: "192.0.2.1", IPVersion: ipVersionV4},
	}); err != nil {
		t.Fatalf("DeleteLiveRouteBatch: %v", err)
	}
	stats, err = db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertCounts(t, "after delete", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
		ipv4Routes: 1, ipv6Routes: 1, liveRoutes: 2,
	})
}

// TestLiveCountsSeededFromDatabaseAtStartup writes rows, reopens the same
// database file, and checks the counts and the prefix distribution come back
// from the seed scan rather than starting at zero.
func TestLiveCountsSeededFromDatabaseAtStartup(t *testing.T) {
	cfg := &config.Config{StateDir: t.TempDir()}

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

	ts := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	if err := db.GetOrCreateASNBatch(map[int]time.Time{64500: ts, 64501: ts, 64502: ts}); err != nil {
		t.Fatalf("GetOrCreateASNBatch: %v", err)
	}
	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, ts),
		mkV6Route("2001:db8::/32", 64502, ts),
	}); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}
	if err := db.Close(); err != nil {
		t.Fatalf("Close: %v", err)
	}

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

	stats, err := reopened.GetStatsContext(context.Background())
	if err != nil {
		t.Fatalf("GetStatsContext after reopen: %v", err)
	}
	if stats.ASNs != 3 {
		t.Errorf("seeded ASNs = %d, want 3", stats.ASNs)
	}
	if stats.IPv4Routes != 1 || stats.IPv6Routes != 1 || stats.LiveRoutes != 2 {
		t.Errorf("seeded routes = (v4 %d, v6 %d, total %d), want (1, 1, 2)",
			stats.IPv4Routes, stats.IPv6Routes, stats.LiveRoutes)
	}
	assertDistribution(t, "seeded IPv4 distribution", stats.IPv4PrefixDistribution,
		[]PrefixDistribution{{MaskLength: 24, Count: 1}})
	assertDistribution(t, "seeded IPv6 distribution", stats.IPv6PrefixDistribution,
		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
}

// TestPrefixDistributionTracksWrites checks that the prefix distribution the
// stats read reports stays exact across each kind of write, through both the
// batch and the single-route methods, and that after every step it equals what
// the distribution query reads from the route tables.
func TestPrefixDistributionTracksWrites(t *testing.T) {
	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)
	shared := mkV4Route(t, "198.51.100.0/24", 64500, ts)
	sharedSecondPeer := mkV4Route(t, "198.51.100.0/24", 64500, ts)
	sharedSecondPeer.PeerIP = "192.0.2.2"
	other := mkV4Route(t, "203.0.113.0/24", 64501, ts)
	wide := mkV4Route(t, "172.16.0.0/16", 64502, ts)
	v6 := mkV6Route("2001:db8::/32", 64503, ts)

	all := []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 2}}
	v6Only := []PrefixDistribution{{MaskLength: 32, Count: 1}}

	steps := []struct {
		name   string
		write  func() error
		wantV4 []PrefixDistribution
		wantV6 []PrefixDistribution
	}{
		{
			name:  "empty database",
			write: func() error { return nil },
		},
		{
			name:   "new prefixes",
			write:  func() error { return db.UpsertLiveRouteBatch([]*LiveRoute{shared, other, wide, v6}) },
			wantV4: all, wantV6: v6Only,
		},
		{
			name:   "re-announcement",
			write:  func() error { return db.UpsertLiveRouteBatch([]*LiveRoute{shared, other, wide, v6}) },
			wantV4: all, wantV6: v6Only,
		},
		{
			name:   "second peer announces a prefix that already has a route",
			write:  func() error { return db.UpsertLiveRoute(sharedSecondPeer) },
			wantV4: all, wantV6: v6Only,
		},
		{
			name: "withdrawal of a route that is not the last for its prefix",
			write: func() error {
				return db.DeleteLiveRouteBatch([]LiveRouteDeletion{
					{Prefix: shared.Prefix, OriginASN: shared.OriginASN, PeerIP: shared.PeerIP, IPVersion: ipVersionV4},
				})
			},
			wantV4: all, wantV6: v6Only,
		},
		{
			name: "withdrawal of the last route for a prefix",
			write: func() error {
				return db.DeleteLiveRoute(sharedSecondPeer.Prefix, sharedSecondPeer.OriginASN, sharedSecondPeer.PeerIP)
			},
			wantV4: []PrefixDistribution{{MaskLength: 16, Count: 1}, {MaskLength: 24, Count: 1}},
			wantV6: v6Only,
		},
		{
			name: "withdrawal of every remaining route, one without an origin ASN",
			write: func() error {
				return db.DeleteLiveRouteBatch([]LiveRouteDeletion{
					{Prefix: other.Prefix, PeerIP: other.PeerIP, IPVersion: ipVersionV4},
					{Prefix: wide.Prefix, OriginASN: wide.OriginASN, PeerIP: wide.PeerIP, IPVersion: ipVersionV4},
					{Prefix: v6.Prefix, OriginASN: v6.OriginASN, PeerIP: v6.PeerIP, IPVersion: ipVersionV6},
				})
			},
		},
	}

	ctx := context.Background()
	for _, step := range steps {
		if err := step.write(); err != nil {
			t.Fatalf("%s: %v", step.name, err)
		}

		stats, err := db.GetStatsContext(ctx)
		if err != nil {
			t.Fatalf("%s: GetStatsContext: %v", step.name, err)
		}
		queryV4, queryV6, err := db.GetPrefixDistributionContext(ctx)
		if err != nil {
			t.Fatalf("%s: GetPrefixDistributionContext: %v", step.name, err)
		}

		assertDistribution(t, step.name+": IPv4 distribution", stats.IPv4PrefixDistribution, step.wantV4)
		assertDistribution(t, step.name+": IPv6 distribution", stats.IPv6PrefixDistribution, step.wantV6)
		assertDistribution(t, step.name+": IPv4 distribution query", queryV4, step.wantV4)
		assertDistribution(t, step.name+": IPv6 distribution query", queryV6, step.wantV6)
	}
}

// TestStatsDistributionNeedsNoQuery checks that the stats read serves the prefix
// distribution from memory. The read's deadline has already passed, so any
// query it ran would fail at once. The distribution query it used to run read
// every live route and passed the /api/v1/stats deadline on a large database
// (issue 30).
func TestStatsDistributionNeedsNoQuery(t *testing.T) {
	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)
	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, ts),
		mkV6Route("2001:db8::/32", 64501, ts),
	}); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

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

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	assertDistribution(t, "IPv4 distribution", stats.IPv4PrefixDistribution,
		[]PrefixDistribution{{MaskLength: 24, Count: 1}})
	assertDistribution(t, "IPv6 distribution", stats.IPv6PrefixDistribution,
		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
}

// TestStatsRouteTimestamps checks the oldest/newest route timestamps are read
// from the right rows across both tables and parse into time.Time. The old
// MIN/MAX union query read its result into *time.Time, which the driver could
// not parse, so it logged a warning every call and left both timestamps nil.
func TestStatsRouteTimestamps(t *testing.T) {
	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()

	empty, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext on empty database: %v", err)
	}
	if empty.OldestRoute != nil || empty.NewestRoute != nil {
		t.Fatalf("empty database timestamps = (%v, %v), want (nil, nil)",
			empty.OldestRoute, empty.NewestRoute)
	}

	base := time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC)
	oldest := base
	newest := base.Add(2 * time.Minute)
	if err := db.UpsertLiveRouteBatch([]*LiveRoute{
		mkV4Route(t, "198.51.100.0/24", 64500, base.Add(time.Minute)),
		mkV4Route(t, "203.0.113.0/24", 64501, oldest),
		mkV6Route("2001:db8::/32", 64502, newest),
	}); err != nil {
		t.Fatalf("UpsertLiveRouteBatch: %v", err)
	}

	stats, err := db.GetStatsContext(ctx)
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	if stats.OldestRoute == nil || !stats.OldestRoute.Equal(oldest) {
		t.Errorf("OldestRoute = %v, want %v", stats.OldestRoute, oldest)
	}
	if stats.NewestRoute == nil || !stats.NewestRoute.Equal(newest) {
		t.Errorf("NewestRoute = %v, want %v", stats.NewestRoute, newest)
	}
}

// TestLiveCountsConcurrentReadWrite runs writers and stats readers at once so
// the race detector proves the counters are safe under concurrent use.
func TestLiveCountsConcurrentReadWrite(t *testing.T) {
	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)

	const writers = 4
	var wg sync.WaitGroup

	wg.Add(writers)
	for w := range writers {
		go func(base int) {
			defer wg.Done()
			for i := range 25 {
				asn := 65000 + base*100 + i
				route := mkV6Route("2001:db8::/32", asn, ts)
				route.PeerIP = "2001:db8::" + uuid.NewString()
				if err := db.UpsertLiveRoute(route); err != nil {
					t.Errorf("UpsertLiveRoute: %v", err)

					return
				}
			}
		}(w)
	}

	var readerWG sync.WaitGroup
	readerWG.Add(1)
	stop := make(chan struct{})
	go func() {
		defer readerWG.Done()
		for {
			select {
			case <-stop:
				return
			default:
				if _, err := db.GetStatsContext(context.Background()); err != nil {
					t.Errorf("GetStatsContext: %v", err)

					return
				}
			}
		}
	}()

	wg.Wait()
	close(stop)
	readerWG.Wait()

	stats, err := db.GetStatsContext(context.Background())
	if err != nil {
		t.Fatalf("GetStatsContext: %v", err)
	}
	if want := writers * 25; stats.IPv6Routes != want {
		t.Errorf("IPv6Routes = %d, want %d", stats.IPv6Routes, want)
	}
	// Every writer announced the same prefix, so it counts once.
	assertDistribution(t, "IPv6 distribution", stats.IPv6PrefixDistribution,
		[]PrefixDistribution{{MaskLength: 32, Count: 1}})
}

type wantCounts struct {
	asns       int
	prefixes   int
	peerings   int
	peers      int
	ipv4Routes int
	ipv6Routes int
	liveRoutes int
}

func assertCounts(t *testing.T, when string, got Stats, want wantCounts) {
	t.Helper()

	if got.ASNs != want.asns {
		t.Errorf("%s: ASNs = %d, want %d", when, got.ASNs, want.asns)
	}
	if got.Prefixes != want.prefixes {
		t.Errorf("%s: Prefixes = %d, want %d", when, got.Prefixes, want.prefixes)
	}
	if got.Peerings != want.peerings {
		t.Errorf("%s: Peerings = %d, want %d", when, got.Peerings, want.peerings)
	}
	if got.Peers != want.peers {
		t.Errorf("%s: Peers = %d, want %d", when, got.Peers, want.peers)
	}
	if got.IPv4Routes != want.ipv4Routes {
		t.Errorf("%s: IPv4Routes = %d, want %d", when, got.IPv4Routes, want.ipv4Routes)
	}
	if got.IPv6Routes != want.ipv6Routes {
		t.Errorf("%s: IPv6Routes = %d, want %d", when, got.IPv6Routes, want.ipv6Routes)
	}
	if got.LiveRoutes != want.liveRoutes {
		t.Errorf("%s: LiveRoutes = %d, want %d", when, got.LiveRoutes, want.liveRoutes)
	}
}

func assertDistribution(t *testing.T, what string, got, want []PrefixDistribution) {
	t.Helper()

	if !slices.Equal(got, want) {
		t.Errorf("%s = %v, want %v", what, got, want)
	}
}