EvoBGP/internal/repository/telemetry_integration_test.go

package repository

import (
	"context"
	"errors"
	"testing"
	"time"

	"evobgp/internal/birdfmt"
	"evobgp/internal/store"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
)

func TestTelemetryLeaseSnapshotAndDiscoveryIntegration(t *testing.T) {
	p, tenant := queueIntegration(t)
	ctx := context.Background()
	speaker := uuid.NewString()
	_, err := p.pool.Exec(ctx, `INSERT INTO bgp_speaker(id,tenant_id,role,meta_json) VALUES($1,$2,'replica',$3::jsonb)`, speaker, tenant, `{"agent_domain":"node.example","agent_secret":"secret"}`)
	if err != nil {
		t.Fatal(err)
	}
	first, err := p.ClaimTelemetry(ctx, "owner-a", 1, time.Minute, "")
	if err != nil || len(first) != 1 || first[0].SpeakerID != speaker {
		t.Fatalf("first claim=%+v %v", first, err)
	}
	if second, err := p.ClaimTelemetry(ctx, "owner-b", 1, time.Minute, ""); err != nil || len(second) != 0 {
		t.Fatalf("overlapping claim=%+v %v", second, err)
	}
	at := time.Now().UTC()
	result := store.TelemetryResult{BGP: store.TelemetryBGP{TelemetryObservation: store.TelemetryObservation{Source: "agent_bird_protocols", ObservedAt: at, ReceivedAt: at}, Sessions: []birdfmt.BGPSession{{Name: birdfmt.DynamicPeerProtocolPrefix + "test", Neighbor: "192.0.2.8", NeighborAS: 64512, State: "Established"}}}, Health: store.TelemetryHealth{TelemetryObservation: store.TelemetryObservation{Source: "agent_health", ObservedAt: at, ReceivedAt: at}, OK: true}}
	if err := p.CompleteTelemetry(ctx, first[0], result, at.Add(time.Hour)); err != nil {
		t.Fatal(err)
	}
	other := &Postgres{pool: p.pool}
	snaps, err := other.ListTelemetry(ctx, tenant)
	if err != nil || len(snaps) != 1 || len(snaps[0].BGP.Sessions) != 1 || snaps[0].ConfigFingerprint == "" {
		t.Fatalf("restarted repository snapshot=%+v %v", snaps, err)
	}
	discoveries, err := other.ListPeerDiscoveries(tenant, store.PeerDiscoveryPending)
	if err != nil || len(discoveries) != 1 {
		t.Fatalf("fenced discovery=%+v %v", discoveries, err)
	}
	if requested, err := other.RequestTelemetryRefresh(ctx, tenant, time.Minute); err != nil || !requested {
		t.Fatalf("refresh=%v %v", requested, err)
	}
	if requested, err := other.RequestTelemetryRefresh(ctx, tenant, time.Minute); err != nil || requested {
		t.Fatalf("refresh not coalesced=%v %v", requested, err)
	}
	newClaim, err := other.ClaimTelemetry(ctx, "owner-b", 1, time.Minute, "")
	if err != nil || len(newClaim) != 1 || newClaim[0].FencingToken <= first[0].FencingToken {
		t.Fatalf("new lease=%+v %v", newClaim, err)
	}
	if err := p.CompleteTelemetry(ctx, first[0], result, at.Add(time.Hour)); !errors.Is(err, store.ErrTelemetryLeaseLost) {
		t.Fatalf("late result accepted: %v", err)
	}

	// Expire the lease inside the discovery write, after the initial fence check.
	// The final fence must roll back BOTH snapshot and discovery atomically.
	function := pgx.Identifier{"test_telemetry_expiry_" + uuid.NewString()}.Sanitize()
	trigger := pgx.Identifier{"test_telemetry_expiry_" + uuid.NewString()}.Sanitize()
	_, err = p.pool.Exec(ctx, `CREATE FUNCTION `+function+`() RETURNS trigger LANGUAGE plpgsql AS $$
	BEGIN
	 IF NEW.tenant_id::text = TG_ARGV[0] THEN
	  UPDATE telemetry_snapshot SET lease_until='epoch'::timestamptz
	  WHERE tenant_id=NEW.tenant_id AND target_id=TG_ARGV[1];
	 END IF;
	 RETURN NEW;
	END $$`)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _, _ = p.pool.Exec(context.Background(), `DROP FUNCTION IF EXISTS `+function+`() CASCADE`) })
	// Both arguments are generated UUIDs, not untrusted SQL input.
	_, err = p.pool.Exec(ctx, `CREATE TRIGGER `+trigger+` BEFORE INSERT ON bgp_peer_discovery FOR EACH ROW EXECUTE FUNCTION `+function+`('`+tenant+`','`+speaker+`')`)
	if err != nil {
		t.Fatal(err)
	}
	result.BGP.Sessions = []birdfmt.BGPSession{{Name: birdfmt.DynamicPeerProtocolPrefix + "rollback", Neighbor: "192.0.2.9", NeighborAS: 64512, State: "Established"}}
	if err = other.CompleteTelemetry(ctx, newClaim[0], result, at.Add(time.Hour)); !errors.Is(err, store.ErrTelemetryLeaseLost) {
		t.Fatalf("completion after lease expired during discovery: %v", err)
	}
	snaps, err = p.ListTelemetry(ctx, tenant)
	if err != nil || len(snaps) != 1 || len(snaps[0].BGP.Sessions) != 1 || snaps[0].BGP.Sessions[0].Neighbor != "192.0.2.8" {
		t.Fatalf("expired completion changed snapshot: %+v %v", snaps, err)
	}
	discoveries, err = p.ListPeerDiscoveries(tenant, store.PeerDiscoveryPending)
	if err != nil || len(discoveries) != 1 {
		t.Fatalf("expired completion changed discoveries: %+v %v", discoveries, err)
	}
	neighbor, ok := store.NormalizePeerNeighborString(discoveries[0].Neighbor)
	if !ok || neighbor != "192.0.2.8" {
		t.Fatalf("expired completion inserted discovery: %+v", discoveries)
	}
}