EvoBGP/internal/pipeline/cdn_snapshot_read_test.go

package pipeline

import (
	"context"
	"sync/atomic"
	"testing"
	"time"

	"evobgp/internal/store"
)

type snapCalls struct {
	store.Backend
	full   atomic.Int32
	narrow atomic.Int32
}

func (s *snapCalls) GetModulePrefixSnapshot(ctx context.Context, tenantID, moduleID string) (*store.ModulePrefixSnapshot, bool, error) {
	s.full.Add(1)
	return s.Backend.GetModulePrefixSnapshot(ctx, tenantID, moduleID)
}

func (s *snapCalls) ListModuleSnapshotSourceRows(ctx context.Context, tenantID, moduleID, source string) ([]store.PrefixRow, error) {
	s.narrow.Add(1)
	return s.Backend.ListModuleSnapshotSourceRows(ctx, tenantID, moduleID, source)
}

func TestNarrowSnapshotReadCheaper(t *testing.T) {
	if !narrowSnapshotReadCheaper(1, 8) {
		t.Fatal("one skipped source should use a narrow read")
	}
	if !narrowSnapshotReadCheaper(3, 8) {
		t.Fatal("a minority of skipped sources should use narrow reads")
	}
	if narrowSnapshotReadCheaper(6, 8) {
		t.Fatal("a majority of skipped sources should share one snapshot read")
	}
	if narrowSnapshotReadCheaper(0, 4) {
		t.Fatal("nothing to read")
	}
}

func TestRefreshModuleIngestDoesNotRereadSnapshot(t *testing.T) {
	t.Setenv("EVOBGP_CDN_ALLOW_PRIVATE", "1")
	m := store.NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()
	mod, err := m.CreateModule(tenant, &store.Module{Type: "CDN_CIDRS", Name: "cdn-once", Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	interval := 3600
	fresh := time.Now().UTC().Add(-30 * time.Second)
	src, err := m.CreateCDNSource(tenant, mod.ID, &store.CDNSource{
		SourceKind:         "txt",
		URL:                "https://cdn.example/prefixes.txt",
		RefreshIntervalSec: &interval,
		LastRefreshedAt:    &fresh,
	})
	if err != nil {
		t.Fatal(err)
	}
	rows := []store.PrefixRow{{Prefix: "203.0.113.0/24", Source: cdnSourceKey(src.ID)}}
	if err := m.SetModulePrefixSnapshot(context.Background(), tenant, mod.ID, "hash", rows); err != nil {
		t.Fatal(err)
	}
	counted := &snapCalls{Backend: m}
	if err := RefreshModuleIngest(context.Background(), counted, nil, tenant, mod.ID); err != nil {
		t.Fatal(err)
	}
	if counted.full.Load() != 1 || counted.narrow.Load() != 0 {
		t.Fatalf("snapshot reads full=%d narrow=%d", counted.full.Load(), counted.narrow.Load())
	}
}

func TestSkippedCDNSourceUsesNarrowReadWithoutPrior(t *testing.T) {
	m := store.NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()
	mod, err := m.CreateModule(tenant, &store.Module{Type: "CDN_CIDRS", Name: "cdn-narrow", Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	interval := 3600
	fresh := time.Now().UTC().Add(-30 * time.Second)
	src := &store.CDNSource{ID: "only", SourceKind: "txt", URL: "https://cdn.example/one", RefreshIntervalSec: &interval, LastRefreshedAt: &fresh}
	want := []store.PrefixRow{{Prefix: "203.0.113.8/32", Source: "cdn:only"}}
	if err := m.SetModulePrefixSnapshot(context.Background(), tenant, mod.ID, "hash", append([]store.PrefixRow{{Prefix: "198.51.100.0/24", Source: "cdn:other"}}, want...)); err != nil {
		t.Fatal(err)
	}
	counted := &snapCalls{Backend: m}
	got, err := collectCDNPrefixRows(context.Background(), counted, nil, tenant, mod, []*store.CDNSource{src}, modulePrior{})
	if err != nil {
		t.Fatal(err)
	}
	if counted.full.Load() != 0 || counted.narrow.Load() != 1 {
		t.Fatalf("snapshot reads full=%d narrow=%d", counted.full.Load(), counted.narrow.Load())
	}
	if len(got) != 1 || got[0].Prefix != "203.0.113.8/32" || got[0].Source != "cdn:only" {
		t.Fatalf("rows=%+v", got)
	}
}

func TestManySkippedCDNSourcesShareOneSnapshotRead(t *testing.T) {
	m := store.NewMemory()
	m.SeedDemo()
	tenant, _, _, _, _ := m.DemoIDs()
	mod, err := m.CreateModule(tenant, &store.Module{Type: "CDN_CIDRS", Name: "cdn-shared", Enabled: true})
	if err != nil {
		t.Fatal(err)
	}
	interval := 3600
	fresh := time.Now().UTC().Add(-30 * time.Second)
	sources := []*store.CDNSource{
		{ID: "a", URL: "https://cdn.example/a", RefreshIntervalSec: &interval, LastRefreshedAt: &fresh},
		{ID: "b", URL: "https://cdn.example/b", RefreshIntervalSec: &interval, LastRefreshedAt: &fresh},
	}
	stored := []store.PrefixRow{
		{Prefix: "203.0.113.1/32", Source: "cdn:a"},
		{Prefix: "203.0.113.2/32", Source: "cdn:b"},
		{Prefix: "198.51.100.0/24", Source: "cdn:gone"},
	}
	if err := m.SetModulePrefixSnapshot(context.Background(), tenant, mod.ID, "hash", stored); err != nil {
		t.Fatal(err)
	}
	counted := &snapCalls{Backend: m}
	got, err := collectCDNPrefixRows(context.Background(), counted, nil, tenant, mod, sources, modulePrior{})
	if err != nil {
		t.Fatal(err)
	}
	if counted.full.Load() != 1 || counted.narrow.Load() != 0 {
		t.Fatalf("snapshot reads full=%d narrow=%d", counted.full.Load(), counted.narrow.Load())
	}
	seen := map[string]bool{}
	for _, row := range got {
		seen[row.Source] = true
	}
	if !seen["cdn:a"] || !seen["cdn:b"] || seen["cdn:gone"] {
		t.Fatalf("rows=%+v", got)
	}
}