flock/pkg/agent/state_test.go

package agent

import (
	"os"
	"path/filepath"
	"testing"
	"time"
)

func newStore(t *testing.T) (*Store, string) {
	t.Helper()
	dir := t.TempDir()
	path := filepath.Join(dir, "allocations.json")
	s, err := NewStore(path, "host001")
	if err != nil {
		t.Fatalf("NewStore: %v", err)
	}
	return s, path
}

func TestStore_EmptyOnFirstOpen(t *testing.T) {
	s, _ := newStore(t)
	if got := len(s.Snapshot()); got != 0 {
		t.Fatalf("Snapshot len = %d, want 0", got)
	}
}

func TestStore_UpsertGetDelete(t *testing.T) {
	s, path := newStore(t)
	a := Allocation{
		ContainerID: "abc",
		Namespace:   "mail",
		PodName:     "stalwart-0",
		OwnerUID:    "uid-1",
		IP6:         "2602:817:3000:f001::1",
		State:       StateCommitted,
		AllocatedAt: time.Now().UTC().Truncate(time.Second),
	}
	if err := s.Upsert(a); err != nil {
		t.Fatalf("Upsert: %v", err)
	}

	got, ok := s.Get("abc")
	if !ok || got.PodName != "stalwart-0" {
		t.Fatalf("Get after Upsert: ok=%v got=%+v", ok, got)
	}

	// Round-trip: a fresh Store reading the same path sees the entry.
	s2, err := NewStore(path, "host001")
	if err != nil {
		t.Fatalf("reopen: %v", err)
	}
	if got, ok := s2.Get("abc"); !ok || got.IP6 != a.IP6 {
		t.Fatalf("reopen Get: ok=%v got=%+v", ok, got)
	}

	if err := s.Delete("abc"); err != nil {
		t.Fatalf("Delete: %v", err)
	}
	if _, ok := s.Get("abc"); ok {
		t.Fatalf("entry still present after Delete")
	}
}

func TestStore_UpsertReplacesByContainerID(t *testing.T) {
	s, _ := newStore(t)
	must := func(err error) {
		t.Helper()
		if err != nil {
			t.Fatal(err)
		}
	}
	must(s.Upsert(Allocation{ContainerID: "abc", IP6: "::1", State: StatePending}))
	must(s.Upsert(Allocation{ContainerID: "abc", IP6: "::2", State: StateCommitted}))
	if got := len(s.Snapshot()); got != 1 {
		t.Fatalf("len = %d, want 1 (Upsert should replace)", got)
	}
	if a, _ := s.Get("abc"); a.IP6 != "::2" || a.State != StateCommitted {
		t.Fatalf("Upsert did not replace: %+v", a)
	}
}

func TestStore_PendingContainerIDs(t *testing.T) {
	s, _ := newStore(t)
	_ = s.Upsert(Allocation{ContainerID: "p1", State: StatePending})
	_ = s.Upsert(Allocation{ContainerID: "c1", State: StateCommitted})
	_ = s.Upsert(Allocation{ContainerID: "p2", State: StatePending})

	pend := s.PendingContainerIDs()
	if len(pend) != 2 {
		t.Fatalf("PendingContainerIDs len = %d, want 2", len(pend))
	}
	have := map[string]bool{pend[0]: true, pend[1]: true}
	if !have["p1"] || !have["p2"] {
		t.Fatalf("PendingContainerIDs = %v, want p1,p2", pend)
	}
}

func TestStore_RejectsWrongVersion(t *testing.T) {
	dir := t.TempDir()
	path := filepath.Join(dir, "allocations.json")
	if err := os.WriteFile(path, []byte(`{"version":99,"node":"x","allocations":[]}`), 0o600); err != nil {
		t.Fatal(err)
	}
	if _, err := NewStore(path, "x"); err == nil {
		t.Fatalf("expected error on bad version, got nil")
	}
}

func TestStore_AtomicWriteDurability(t *testing.T) {
	// We can't simulate a real power-loss in unit tests, but we can verify
	// that no .tmp file is left behind after a successful flush, and that
	// the rename target is intact.
	s, path := newStore(t)
	if err := s.Upsert(Allocation{ContainerID: "x", State: StateCommitted}); err != nil {
		t.Fatal(err)
	}
	if _, err := os.Stat(path + ".tmp"); !os.IsNotExist(err) {
		t.Fatalf(".tmp leaked: err=%v", err)
	}
	b, err := os.ReadFile(path)
	if err != nil || len(b) == 0 {
		t.Fatalf("final file unreadable: err=%v len=%d", err, len(b))
	}
}