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

package db

import (
	"context"
	"fmt"
	"testing"
	"time"
)

func validatorState(t *testing.T, d *DB, id string) *Validator {
	t.Helper()
	v, err := d.GetValidator(testCtx(t), id)
	if err != nil {
		t.Fatalf("get validator %s: %v", id, err)
	}
	return v
}

func testCtx(t *testing.T) context.Context {
	t.Helper()
	return context.Background()
}

// claimFor2 seeds an address and returns its id without claiming it.
func claimFor2(t *testing.T, d *DB, addr string) int64 {
	t.Helper()
	ctx := testCtx(t)
	if err := d.SeedQueue(ctx, []string{addr}); err != nil {
		t.Fatalf("seed %s: %v", addr, err)
	}
	ip, err := d.GetIPByAddress(ctx, addr)
	if err != nil {
		t.Fatalf("get %s: %v", addr, err)
	}
	return ip.ID
}

// claimFor seeds an address and claims it for the validator.
func claimFor(t *testing.T, d *DB, addr, validatorID string) *IPQueueItem {
	t.Helper()
	ctx := testCtx(t)
	if err := d.SeedQueue(ctx, []string{addr}); err != nil {
		t.Fatalf("seed %s: %v", addr, err)
	}
	item, err := d.ClaimNextQueued(ctx, validatorID, time.Minute)
	if err != nil || item == nil {
		t.Fatalf("claim %s for %s: item=%v err=%v", addr, validatorID, item, err)
	}
	return item
}

func TestHeartbeatKeepsAssignedWhenValidatorStillHoldsAnAddress(t *testing.T) {
	d, ctx := newTestDB(t)
	_ = d.AdminCreateValidator(ctx, "v1", "p1")
	_ = d.AdminCreateValidator(ctx, "v2", "p2")
	item := claimFor(t, d, "1.1.1.1", "v1")
	_ = d.MarkValidatorUnreachable(ctx, "v1")
	_ = d.MarkValidatorUnreachable(ctx, "v2")

	if err := d.Heartbeat(ctx, "v1"); err != nil {
		t.Fatal(err)
	}
	if v := validatorState(t, d, "v1"); v.State != ValidatorAssigned || v.CurrentIPID == nil || *v.CurrentIPID != item.ID {
		t.Fatalf("v1 after heartbeat: state=%s current_ip=%v, want assigned to %d", v.State, v.CurrentIPID, item.ID)
	}
	if err := d.Heartbeat(ctx, "v2"); err != nil {
		t.Fatal(err)
	}
	if v := validatorState(t, d, "v2"); v.State != ValidatorIdle {
		t.Fatalf("v2 after heartbeat: state=%s, want idle", v.State)
	}
}

func TestRegisterValidatorReactivationKeepsAssignedAddress(t *testing.T) {
	d, ctx := newTestDB(t)
	_ = d.AdminCreateValidator(ctx, "v1", "p1")
	item := claimFor(t, d, "1.1.1.1", "v1")
	_ = d.MarkValidatorUnreachable(ctx, "v1")

	if err := d.RegisterValidator(ctx, "v1", "host", "p1", "v"); err != nil {
		t.Fatal(err)
	}
	if v := validatorState(t, d, "v1"); v.State != ValidatorAssigned || v.CurrentIPID == nil || *v.CurrentIPID != item.ID {
		t.Fatalf("after re-register: state=%s current_ip=%v, want assigned to %d", v.State, v.CurrentIPID, item.ID)
	}
}

func TestStaleReleasesLeaveTheCurrentAddressAlone(t *testing.T) {
	cases := map[string]func(d *DB, ipID int64) error{
		"ReleaseFIP":      func(d *DB, id int64) error { return d.ReleaseFIP(context.Background(), id, "v1") },
		"RequeueOrFail":   func(d *DB, id int64) error { return d.RequeueOrFail(context.Background(), id, "v1", 3) },
		"MarkFIPOccupied": func(d *DB, id int64) error { return d.MarkFIPOccupied(context.Background(), id, "v1") },
		"FreeValidator":   func(d *DB, id int64) error { return d.FreeValidator(context.Background(), "v1", id) },
	}
	for name, release := range cases {
		t.Run(name, func(t *testing.T) {
			d, ctx := newTestDB(t)
			_ = d.AdminCreateValidator(ctx, "v1", "p1")
			stale := claimFor(t, d, "1.1.1.1", "v1")
			// v1 has moved on to another address (set directly: the claim path
			// itself refuses a validator that is still busy).
			cur := claimFor2(t, d, "2.2.2.2")
			if _, err := d.ExecContext(ctx, `UPDATE ip_queue SET state='checking', owner_validator_id='v1' WHERE id=?`, cur); err != nil {
				t.Fatal(err)
			}
			if _, err := d.ExecContext(ctx, `UPDATE validators SET state='assigned', current_ip_id=? WHERE validator_id='v1'`, cur); err != nil {
				t.Fatal(err)
			}

			if err := release(d, stale.ID); err != nil {
				t.Fatalf("%s: %v", name, err)
			}
			v := validatorState(t, d, "v1")
			if v.State != ValidatorAssigned || v.CurrentIPID == nil || *v.CurrentIPID != cur {
				t.Fatalf("%s freed a validator that holds another address: state=%s current_ip=%v", name, v.State, v.CurrentIPID)
			}
		})
	}
}

// The right address frees the validator, but an unreachable one stays so.
func TestReleaseKeepsUnreachableValidatorUnreachable(t *testing.T) {
	d, ctx := newTestDB(t)
	_ = d.AdminCreateValidator(ctx, "v1", "p1")
	_ = d.AdminCreateValidator(ctx, "v2", "p2")
	a := claimFor(t, d, "1.1.1.1", "v1")
	b := claimFor(t, d, "2.2.2.2", "v2")
	_ = d.MarkValidatorUnreachable(ctx, "v1")

	if err := d.ReleaseFIP(ctx, a.ID, "v1"); err != nil {
		t.Fatal(err)
	}
	if v := validatorState(t, d, "v1"); v.State != ValidatorUnreachable || v.CurrentIPID != nil {
		t.Fatalf("v1: state=%s current_ip=%v, want unreachable and empty", v.State, v.CurrentIPID)
	}
	if err := d.RequeueOrFail(ctx, b.ID, "v2", 3); err != nil {
		t.Fatal(err)
	}
	if v := validatorState(t, d, "v2"); v.State != ValidatorIdle || v.CurrentIPID != nil {
		t.Fatalf("v2: state=%s current_ip=%v, want idle and empty", v.State, v.CurrentIPID)
	}
}

func TestClaimRefusesValidatorThatStillHoldsAnAddress(t *testing.T) {
	d, ctx := newTestDB(t)
	_ = d.AdminCreateValidator(ctx, "v1", "p1")
	first := claimFor(t, d, "1.1.1.1", "v1")
	// An old version could leave a validator idle while it still pointed at an address.
	if _, err := d.ExecContext(ctx, `UPDATE validators SET state='idle' WHERE validator_id='v1'`); err != nil {
		t.Fatal(err)
	}
	_ = d.SeedQueue(ctx, []string{"2.2.2.2"})
	got, err := d.ClaimNextQueued(ctx, "v1", time.Minute)
	if err != nil || got != nil {
		t.Fatalf("claim for a validator that holds %d: item=%v err=%v, want nothing", first.ID, got, err)
	}
	if ip, _ := d.GetIPByAddress(ctx, "2.2.2.2"); ip.State != IPQueued {
		t.Fatalf("2.2.2.2 is %s, want queued", ip.State)
	}
}

func TestListFIPRefsSkipsFinishedAddresses(t *testing.T) {
	d, ctx := newTestDB(t)
	var addrs []string
	for i := 0; i < 6; i++ {
		addrs = append(addrs, fmt.Sprintf("10.0.0.%d", i+1))
	}
	_ = d.SeedQueue(ctx, addrs)
	states := []string{"done", "failed", "occupied", "awaiting_self_check", "checking", "aggregating"}
	for i, a := range addrs {
		if _, err := d.ExecContext(ctx, `UPDATE ip_queue SET state=?, fip_id=? WHERE ip_address=?`, states[i], fmt.Sprintf("fip-%d", i), a); err != nil {
			t.Fatal(err)
		}
	}
	refs, err := d.ListFIPRefs(ctx)
	if err != nil {
		t.Fatal(err)
	}
	if len(refs) != 3 {
		t.Fatalf("ListFIPRefs returned %d rows, want the 3 unfinished ones: %+v", len(refs), refs)
	}
	byAddr, err := d.ListFIPRefsByAddresses(ctx, addrs)
	if err != nil || len(byAddr) != 3 {
		t.Fatalf("ListFIPRefsByAddresses returned %d rows (err %v), want 3", len(byAddr), err)
	}
}