routewatch/internal/database/counts_internal_test.go

package database

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"git.eeqj.de/sneak/routewatch/internal/config"
	"git.eeqj.de/sneak/routewatch/internal/logger"
)

// testV4PeerIP is the peer the IPv4 rows in these tests are learned from.
const testV4PeerIP = "192.0.2.1"

// testV6Prefix is the IPv6 prefix of these tests.
const testV6Prefix = "2001:db8::/32"

// TestLiveCountsTrackWritesInRealtime checks that the stats counts start at
// zero and reflect each write the moment it commits (no recompute, no timer).
// 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) {
	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()

	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 {
		t.Fatalf("empty database counts nonzero: %+v", empty)
	}

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

	writeASNsPrefixesAndPeers(t, db, ts)

	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,
	})
}

// writeASNsPrefixesAndPeers writes two ASNs, two prefixes, one peer and one
// peering: the rows the stats count.
func writeASNsPrefixesAndPeers(t *testing.T, db *Database, ts time.Time) {
	t.Helper()

	seen := SeenTimes{First: ts, Last: ts}

	err := db.GetOrCreateASNBatch(map[int]SeenTimes{64500: seen, 64501: seen})
	if err != nil {
		t.Fatalf("GetOrCreateASNBatch: %v", err)
	}

	err = db.UpdatePrefixesBatch(map[string]time.Time{
		"198.51.100.0/24": ts,
		testV6Prefix:      ts,
	})
	if err != nil {
		t.Fatalf("UpdatePrefixesBatch: %v", err)
	}

	err = db.UpdatePeerBatch(map[string]PeerUpdate{
		testV4PeerIP: {
			PeerIP: testV4PeerIP, PeerASN: 64500,
			MessageType: "UPDATE", Seen: seen,
		},
	})
	if err != nil {
		t.Fatalf("UpdatePeerBatch: %v", err)
	}

	err = db.RecordPeering(64500, 64501, seen)
	if err != nil {
		t.Fatalf("RecordPeering: %v", err)
	}
}

// TestLiveCountsSeededFromDatabaseAtStartup writes rows, reopens the same
// database file, and checks the counts come back from the seed scan rather
// than starting at zero.
func TestLiveCountsSeededFromDatabaseAtStartup(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)
	}

	writeASNsPrefixesAndPeers(t, db, time.Date(2026, 1, 2, 3, 4, 5, 0, time.UTC))

	err = db.Close()
	if 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)
	}

	assertCounts(t, "after reopen", stats, wantCounts{
		asns: 2, prefixes: 2, peerings: 1, peers: 1,
	})
}

// 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) {
	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)

	const (
		writers          = 4
		prefixesPerWrite = 25
	)

	var wg sync.WaitGroup
	wg.Add(writers)

	for w := range writers {
		go func(base int) {
			defer wg.Done()

			for i := range prefixesPerWrite {
				prefix := fmt.Sprintf("10.%d.%d.0/24", base, i)

				err := db.UpdatePrefixesBatch(map[string]time.Time{prefix: ts})
				if err != nil {
					t.Errorf("UpdatePrefixesBatch: %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:
				_, err := db.GetStatsContext(context.Background())
				if 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 * prefixesPerWrite; stats.IPv4Prefixes != want {
		t.Errorf("IPv4Prefixes = %d, want %d", stats.IPv4Prefixes, want)
	}
}

type wantCounts struct {
	asns     int
	prefixes int
	peerings int
	peers    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)
	}
}