cloud-ip-validator/internal/db/queries_sites_heartbeat_test.go

package db

import (
	"database/sql"
	"errors"
	"testing"
	"time"
)

func TestRegisterSiteProberUnknownSiteReturnsErrNoRows(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.RegisterSiteProber(ctx, "unknown-site", "host-1"); !errors.Is(err, sql.ErrNoRows) {
		t.Fatalf("expected sql.ErrNoRows for unknown site, got %v", err)
	}
}

func TestRegisterSiteProberSetsHostnameAndIdleState(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}

	sites, err := d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites: %v", err)
	}
	if sites[0].State != SiteUnregistered {
		t.Fatalf("expected freshly-created slot unregistered, got %s", sites[0].State)
	}

	if err := d.RegisterSiteProber(ctx, "site-1", "probe-host-1"); err != nil {
		t.Fatalf("register site prober: %v", err)
	}
	sites, err = d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites after register: %v", err)
	}
	if sites[0].State != SiteIdle {
		t.Fatalf("expected idle after register, got %s", sites[0].State)
	}
	if sites[0].Hostname != "probe-host-1" {
		t.Fatalf("expected hostname set, got %q", sites[0].Hostname)
	}
	if sites[0].LastHeartbeatAt == nil {
		t.Fatalf("expected last_heartbeat_at stamped")
	}
}

func TestRegisterSiteProberReactivatesFromUnreachable(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}
	if err := d.MarkSiteUnreachable(ctx, "site-1"); err != nil {
		t.Fatalf("mark unreachable: %v", err)
	}
	if err := d.RegisterSiteProber(ctx, "site-1", "probe-host-1"); err != nil {
		t.Fatalf("register site prober: %v", err)
	}
	sites, err := d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites: %v", err)
	}
	if sites[0].State != SiteIdle {
		t.Fatalf("expected reactivated to idle, got %s", sites[0].State)
	}
}

func TestSiteHeartbeatRoundTrip(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}
	if err := d.SiteHeartbeat(ctx, "unknown-site"); !errors.Is(err, sql.ErrNoRows) {
		t.Fatalf("expected sql.ErrNoRows for unknown site, got %v", err)
	}

	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
		t.Fatalf("heartbeat: %v", err)
	}
	sites, err := d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites: %v", err)
	}
	// A plain heartbeat (no prior register) doesn't flip an unregistered
	// site to idle — only reactivates from unreachable, mirroring
	// validators' Heartbeat.
	if sites[0].State != SiteUnregistered {
		t.Fatalf("expected still unregistered after bare heartbeat, got %s", sites[0].State)
	}
	if sites[0].LastHeartbeatAt == nil {
		t.Fatalf("expected last_heartbeat_at stamped")
	}
	if sites[0].Hostname != "" {
		t.Fatalf("expected heartbeat not to set hostname, got %q", sites[0].Hostname)
	}

	if err := d.MarkSiteUnreachable(ctx, "site-1"); err != nil {
		t.Fatalf("mark unreachable: %v", err)
	}
	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
		t.Fatalf("heartbeat after unreachable: %v", err)
	}
	sites, err = d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites: %v", err)
	}
	if sites[0].State != SiteIdle {
		t.Fatalf("expected heartbeat to reactivate from unreachable to idle, got %s", sites[0].State)
	}
}

func TestListStaleSiteHeartbeatsAndMarkUnreachable(t *testing.T) {
	d, ctx := newTestDB(t)
	if err := d.UpsertSite(ctx, 1, "site-1"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}
	if err := d.UpsertSite(ctx, 2, "site-2"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}
	// site-3 never heartbeats at all — must not show up as "stale" (it's
	// simply unregistered, which is a distinct, expected state).
	if err := d.UpsertSite(ctx, 3, "site-3"); err != nil {
		t.Fatalf("upsert site: %v", err)
	}

	if err := d.SiteHeartbeat(ctx, "site-1"); err != nil {
		t.Fatalf("heartbeat site-1: %v", err)
	}
	if err := d.SiteHeartbeat(ctx, "site-2"); err != nil {
		t.Fatalf("heartbeat site-2: %v", err)
	}

	cutoff := Now().Add(1 * time.Hour) // everything heartbeated is "before" this
	stale, err := d.ListStaleSiteHeartbeats(ctx, cutoff)
	if err != nil {
		t.Fatalf("list stale site heartbeats: %v", err)
	}
	if len(stale) != 2 {
		t.Fatalf("expected 2 stale sites (site-1, site-2), got %+v", stale)
	}

	for _, s := range stale {
		if err := d.MarkSiteUnreachable(ctx, s.SiteID); err != nil {
			t.Fatalf("mark unreachable %s: %v", s.SiteID, err)
		}
	}

	// Now already-unreachable sites are excluded from a subsequent sweep.
	stale, err = d.ListStaleSiteHeartbeats(ctx, cutoff)
	if err != nil {
		t.Fatalf("list stale site heartbeats after mark: %v", err)
	}
	if len(stale) != 0 {
		t.Fatalf("expected no stale sites left, got %+v", stale)
	}

	sites, err := d.ListSites(ctx)
	if err != nil {
		t.Fatalf("list sites: %v", err)
	}
	byID := map[string]Site{}
	for _, s := range sites {
		byID[s.SiteID] = s
	}
	if byID["site-1"].State != SiteUnreachable || byID["site-2"].State != SiteUnreachable {
		t.Fatalf("expected site-1/site-2 unreachable, got %+v", sites)
	}
	if byID["site-3"].State != SiteUnregistered {
		t.Fatalf("expected site-3 still unregistered (never heartbeated), got %s", byID["site-3"].State)
	}
}