MOPAC/internal/quota/resources_test.go

package quota

import (
	"bufio"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"testing"
	"time"
)

func TestResourceStatsFromFixture(t *testing.T) {
	root := t.TempDir()
	proc := filepath.Join(root, "proc")
	os.MkdirAll(filepath.Join(proc, "pressure"), 0o755)
	os.WriteFile(filepath.Join(proc, "loadavg"), []byte("7.32 5.10 2.20 3/900 12345\n"), 0o644)
	os.WriteFile(filepath.Join(proc, "meminfo"), []byte("MemTotal:       16000000 kB\nMemAvailable:    3145728 kB\n"), 0o644)
	os.WriteFile(filepath.Join(proc, "pressure", "io"), []byte("some avg10=1.00 avg60=2.50 avg300=0.10 total=123456789\nfull avg10=0.00 avg60=0.00 avg300=0.00 total=0\n"), 0o644)
	work := filepath.Join(root, "work")
	os.MkdirAll(work, 0o755)

	st := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
	if st.LoadAvg1m != 7.32 {
		t.Errorf("load = %v, want 7.32", st.LoadAvg1m)
	}
	if st.MemAvailable != 3072 {
		t.Errorf("mem = %v MB, want 3072", st.MemAvailable)
	}
	if !st.HasPSI || st.IODelayPct != 2.5 {
		t.Errorf("io = %v (psi=%v), want 2.5", st.IODelayPct, st.HasPSI)
	}
	if st.DiskFree <= 0 {
		t.Errorf("disk free = %v, want > 0 (statfs on tempdir)", st.DiskFree)
	}

	check := BusyCheck{MaxLoadAvg: 6, MinMemAvailableMB: 4096, MinDiskFreeMB: 1, MaxIODelayPct: 90}
	reasons := check.Evaluate(st)
	joined := strings.Join(reasons, ";")
	for _, want := range []string{"load 7.32 > 6", "mem available 3072MB < 4096MB"} {
		if !strings.Contains(joined, want) {
			t.Errorf("reasons %q missing %q", joined, want)
		}
	}
	// IO under threshold: must NOT appear.
	if strings.Contains(joined, "io delay") {
		t.Errorf("io below threshold must not trip: %q", joined)
	}

	// Missing PSI file: the IO check is skipped, not an error.
	os.Remove(filepath.Join(proc, "pressure", "io"))
	st2 := ReadResourceStats(proc, filepath.Join(root, "sys"), work)
	if st2.HasPSI {
		t.Error("missing PSI must read HasPSI=false")
	}
	if got := (BusyCheck{MaxIODelayPct: 1}).Evaluate(st2); len(got) != 0 {
		t.Errorf("no PSI must never trip IO: %v", got)
	}
}

// TestResourceStatsNeverBusyOnUnreadable mirrors the deploy invariant: read
// failures defer nothing (the gate is advisory, /proc layout changes must
// not stop work).
func TestResourceStatsNeverBusyOnUnreadable(t *testing.T) {
	root := t.TempDir() // empty: no loadavg, no meminfo
	st := ReadResourceStats(filepath.Join(root, "proc"), filepath.Join(root, "sys"), root)
	if got := (BusyCheck{MaxLoadAvg: 1, MinMemAvailableMB: 1, MinDiskFreeMB: 1, MaxIODelayPct: 1}).Evaluate(st); len(got) > 1 {
		t.Errorf("unreadable sources must not trip load/mem (disk may read): %v", got)
	}
}

// fakeRedis is a minimal RESP2 server: enough command surface for the
// SharedState client (SET/GET/INCRBYFLOAT/EXPIRE/SELECT), backed by a map.
type fakeRedis struct {
	mu   sync.Mutex
	data map[string]string
	ttl  map[string]int64
	srv  net.Listener
}

func newFakeRedis(t *testing.T) *fakeRedis {
	t.Helper()
	f := &fakeRedis{data: map[string]string{}, ttl: map[string]int64{}}
	srv, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	f.srv = srv
	go f.serve()
	t.Cleanup(func() { srv.Close() })
	return f
}

func (f *fakeRedis) url(t *testing.T) string {
	t.Helper()
	return "redis://" + f.srv.Addr().String() + "/0"
}

func (f *fakeRedis) serve() {
	for {
		conn, err := f.srv.Accept()
		if err != nil {
			return
		}
		go f.handle(conn)
	}
}

func (f *fakeRedis) handle(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)
	w := bufio.NewWriter(conn)
	for {
		args, err := readCommand(r)
		if err != nil {
			return
		}
		if len(args) == 0 {
			continue
		}
		f.mu.Lock()
		reply := f.exec(args)
		f.mu.Unlock()
		w.WriteString(reply + "\r\n")
		if w.Flush() != nil {
			return
		}
	}
}

func readCommand(r *bufio.Reader) ([]string, error) {
	line, err := r.ReadString('\n')
	if err != nil {
		return nil, err
	}
	line = strings.TrimRight(line, "\r\n")
	if !strings.HasPrefix(line, "*") {
		return strings.Fields(line), nil
	}
	n, _ := strconv.Atoi(line[1:])
	args := make([]string, 0, n)
	for i := 0; i < n; i++ {
		hl, err := r.ReadString('\n')
		if err != nil {
			return nil, err
		}
		hl = strings.TrimRight(hl, "\r\n")
		ln, _ := strconv.Atoi(strings.TrimPrefix(hl, "$"))
		buf := make([]byte, ln+2)
		if _, err := readFull(r, buf); err != nil {
			return nil, err
		}
		args = append(args, string(buf[:ln]))
	}
	return args, nil
}

func readFull(r *bufio.Reader, buf []byte) (int, error) {
	total := 0
	for total < len(buf) {
		n, err := r.Read(buf[total:])
		total += n
		if err != nil {
			return total, err
		}
	}
	return total, nil
}

func (f *fakeRedis) exec(args []string) string {
	cmd := strings.ToUpper(args[0])
	switch cmd {
	case "PING":
		return "+PONG"
	case "SELECT":
		return "+OK"
	case "SET":
		f.data[args[1]] = args[2]
		if len(args) >= 5 && strings.ToUpper(args[3]) == "EX" {
			f.ttl[args[1]], _ = strconv.ParseInt(args[4], 10, 64)
		}
		return "+OK"
	case "GET":
		if v, ok := f.data[args[1]]; ok {
			return fmt.Sprintf("$%d\r\n%s", len(v), v)
		}
		return "$-1"
	case "EXPIRE":
		if _, ok := f.data[args[1]]; ok {
			f.ttl[args[1]], _ = strconv.ParseInt(args[2], 10, 64)
			return ":1"
		}
		return ":0"
	case "INCRBYFLOAT":
		cur := 0.0
		if v, ok := f.data[args[1]]; ok {
			cur, _ = strconv.ParseFloat(v, 64)
		}
		delta, err := strconv.ParseFloat(args[2], 64)
		if err != nil {
			return "-ERR bad delta"
		}
		cur += delta
		f.data[args[1]] = strconv.FormatFloat(cur, 'f', 6, 64)
		return "$" + strconv.Itoa(len(f.data[args[1]])) + "\r\n" + f.data[args[1]]
	default:
		return "-ERR unknown command '" + cmd + "'"
	}
}

func TestSharedStateRoundTrip(t *testing.T) {
	fake := newFakeRedis(t)
	st, err := NewSharedState(fake.url(t))
	if err != nil {
		t.Fatal(err)
	}
	defer st.Close()

	if err := st.setEx("k", "v1", time.Second); err != nil {
		t.Fatal(err)
	}
	if v, err := st.get("k"); err != nil || v != "v1" {
		t.Fatalf("get = %q err=%v", v, err)
	}
	if v, err := st.incrByFloat("ctr", 1.5); err != nil || v != 1.5 {
		t.Fatalf("incr = %v err=%v", v, err)
	}
	if v, err := st.incrByFloat("ctr", 2.25); err != nil || v != 3.75 {
		t.Fatalf("incr2 = %v err=%v (want 3.75)", v, err)
	}
	if v, err := st.get("missing"); err != nil || v != "" {
		t.Fatalf("missing key = %q err=%v (want empty, no error)", v, err)
	}

	// snapshot publish/load through the typed helpers
	snap := &QuotaSnapshot{Account: "zai-9", Source: "usage_url", FetchedAt: time.Now(),
		Buckets: []Bucket{{ID: Bucket5h, Used: 3.5, Limit: 28000}}}
	if err := st.PublishSnapshot(snap, time.Minute); err != nil {
		t.Fatal(err)
	}
	got, err := st.LoadSnapshot("zai-9")
	if err != nil || got == nil {
		t.Fatalf("load = %v err=%v", got, err)
	}
	if b, _ := got.Bucket(Bucket5h); b.Used != 3.5 {
		t.Errorf("bucket round trip: %+v", got.Buckets)
	}
}

func TestSharedStateDownNeverFatal(t *testing.T) {
	// Nothing listens on this port: every op must error softly (nil state
	// ops are no-ops; a dead server returns errors, never panics).
	dead, err := NewSharedState("redis://127.0.0.1:1/0")
	if err != nil {
		t.Fatal(err)
	}
	if err := dead.setEx("k", "v", time.Second); err == nil {
		t.Error("dead redis SET must error")
	}
	if _, err := dead.incrByFloat("k", 1); err == nil {
		t.Error("dead redis INCR must error")
	}
	var nilState *SharedState
	if err := nilState.setEx("k", "v", time.Second); err != nil {
		t.Errorf("nil state must no-op: %v", err)
	}
	if st, err := NewSharedState(""); err != nil || st != nil {
		t.Errorf("empty url must return nil state, got %v %v", st, err)
	}
	if _, err := NewSharedState("redis://"); err == nil {
		t.Error("hostless url must error")
	}
}