at-container-registry/pkg/appview/webhooks/dispatch_quota_test.go

package webhooks

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"net/http/httptest"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"atcr.io/pkg/appview/db"
	"atcr.io/pkg/appview/storage"
	"atcr.io/pkg/atproto"
)

// fakeHold serves io.atcr.hold.getQuota with a configurable usage value so
// the test can drive the edge condition.
type fakeHold struct {
	limit int64
	usage atomic.Int64
}

func (f *fakeHold) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != atproto.HoldGetQuota {
			http.NotFound(w, r)
			return
		}
		_ = json.NewEncoder(w).Encode(map[string]any{
			"userDid":   r.URL.Query().Get("userDid"),
			"totalSize": f.usage.Load(),
			"limit":     f.limit,
		})
	})
}

// fakeReceiver captures webhook deliveries so the test can assert on count.
type fakeReceiver struct {
	mu       sync.Mutex
	payloads [][]byte
	done     chan struct{}
}

func newFakeReceiver(expected int) *fakeReceiver {
	r := &fakeReceiver{done: make(chan struct{}, expected+1)}
	return r
}

func (r *fakeReceiver) handler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		buf := make([]byte, req.ContentLength)
		_, _ = req.Body.Read(buf)
		r.mu.Lock()
		r.payloads = append(r.payloads, buf)
		r.mu.Unlock()
		select {
		case r.done <- struct{}{}:
		default:
		}
		w.WriteHeader(http.StatusOK)
	})
}

func (r *fakeReceiver) count() int {
	r.mu.Lock()
	defer r.mu.Unlock()
	return len(r.payloads)
}

// waitFor blocks up to d for the receiver to accumulate n deliveries.
func (r *fakeReceiver) waitFor(n int, d time.Duration) bool {
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if r.count() >= n {
			return true
		}
		time.Sleep(10 * time.Millisecond)
	}
	return false
}

// TestDispatchForQuotaEdgeTriggered exercises the full dispatcher path against
// a real SQLite DB and a fake hold. Covers:
//   - fires once when usage crosses upward through the threshold
//   - suppresses repeat fires while still above (last_fired_at sticky)
//   - re-arms when usage drops below threshold and fires on the next crossing
//   - non-quota webhooks (push-only) are not picked up by DispatchForQuota
//   - captain / unlimited holds (Limit == nil) never fire
func TestDispatchForQuotaEdgeTriggered(t *testing.T) {
	conn, err := db.InitDB(":memory:", db.LibsqlConfig{})
	if err != nil {
		t.Fatalf("init db: %v", err)
	}
	defer conn.Close()

	const userDID = "did:plc:quotatest"
	if err := db.UpsertUser(conn, &db.User{
		DID: userDID, Handle: "qt.test", PDSEndpoint: "https://pds", LastSeen: time.Now(),
	}); err != nil {
		t.Fatalf("upsert user: %v", err)
	}

	hold := &fakeHold{limit: 1000}
	hold.usage.Store(400) // 40% — below threshold
	holdSrv := httptest.NewServer(hold.handler())
	defer holdSrv.Close()

	receiver := newFakeReceiver(3)
	recvSrv := httptest.NewServer(receiver.handler())
	defer recvSrv.Close()

	// Quota webhook with threshold 75
	quotaHook := &db.Webhook{
		ID:        "wh-quota",
		UserDID:   userDID,
		URL:       recvSrv.URL,
		Triggers:  PackTriggers(TriggerQuota, 75),
		CreatedAt: time.Now().UTC(),
	}
	if err := db.InsertWebhook(conn, quotaHook); err != nil {
		t.Fatalf("insert quota hook: %v", err)
	}

	// Push-only webhook on the same user — must be ignored by DispatchForQuota
	pushHook := &db.Webhook{
		ID:        "wh-push",
		UserDID:   userDID,
		URL:       recvSrv.URL + "/push",
		Triggers:  PackTriggers(TriggerPush, 0),
		CreatedAt: time.Now().UTC(),
	}
	if err := db.InsertWebhook(conn, pushHook); err != nil {
		t.Fatalf("insert push hook: %v", err)
	}

	d := NewDispatcher(conn, atproto.AppviewMetadata{ClientShortName: "ATCR", BaseURL: "https://atcr.test"}, nil)
	// The receiver is an httptest server on loopback, which the delivery guard
	// would otherwise refuse. See ssrf_test.go for the guard's own coverage.
	d.allowLoopbackDeliveryForTest()

	event := storage.QuotaWebhookEvent{
		UserDID:      userDID,
		UserHandle:   "qt.test",
		HoldDID:      holdSrv.URL, // ResolveHoldURL passes URLs through
		HoldEndpoint: holdSrv.URL,
	}

	// 1) Below threshold: no fire
	d.DispatchForQuota(context.Background(), event)
	time.Sleep(50 * time.Millisecond) // give any in-flight goroutine a chance
	if got := receiver.count(); got != 0 {
		t.Fatalf("below threshold: expected 0 deliveries, got %d", got)
	}

	// 2) Cross upward to 80% — should fire exactly once
	hold.usage.Store(800)
	d.DispatchForQuota(context.Background(), event)
	if !receiver.waitFor(1, 2*time.Second) {
		t.Fatalf("crossing upward: webhook did not fire within timeout, count=%d", receiver.count())
	}
	if got := receiver.count(); got != 1 {
		t.Fatalf("crossing upward: expected 1 delivery, got %d", got)
	}

	// Verify payload contents
	var payload QuotaWebhookPayload
	if err := json.Unmarshal(receiver.payloads[0], &payload); err != nil {
		t.Fatalf("unmarshal payload: %v", err)
	}
	if payload.Trigger != "quota" {
		t.Errorf("payload trigger = %q, want %q", payload.Trigger, "quota")
	}
	if payload.QuotaData.UsagePercent != 80 || payload.QuotaData.ThresholdPercent != 75 {
		t.Errorf("payload pcts = (%d, %d), want (80, 75)", payload.QuotaData.UsagePercent, payload.QuotaData.ThresholdPercent)
	}
	if payload.QuotaData.LimitBytes != 1000 || payload.QuotaData.UsageBytes != 800 {
		t.Errorf("payload bytes = (%d / %d), want (800 / 1000)", payload.QuotaData.UsageBytes, payload.QuotaData.LimitBytes)
	}
	if payload.User.DID != userDID {
		t.Errorf("payload user.did = %q, want %q", payload.User.DID, userDID)
	}

	// last_fired_at should be stamped
	stored, err := db.GetWebhookByID(conn, quotaHook.ID)
	if err != nil {
		t.Fatalf("get webhook: %v", err)
	}
	if stored.LastFiredAt == nil {
		t.Fatal("expected last_fired_at to be set after firing")
	}

	// 3) Still above, second push: must not re-fire
	hold.usage.Store(900)
	d.DispatchForQuota(context.Background(), event)
	time.Sleep(50 * time.Millisecond)
	if got := receiver.count(); got != 1 {
		t.Fatalf("still above: expected 1 delivery total, got %d", got)
	}

	// 4) Drop below threshold: re-arms (clears last_fired_at), no fire
	hold.usage.Store(500)
	d.DispatchForQuota(context.Background(), event)
	time.Sleep(50 * time.Millisecond)
	if got := receiver.count(); got != 1 {
		t.Fatalf("drop below: expected 1 delivery total, got %d", got)
	}
	stored, err = db.GetWebhookByID(conn, quotaHook.ID)
	if err != nil {
		t.Fatalf("get webhook after drop: %v", err)
	}
	if stored.LastFiredAt != nil {
		t.Errorf("expected last_fired_at to be cleared after dropping below; got %v", *stored.LastFiredAt)
	}

	// 5) Cross upward again — fires again
	hold.usage.Store(950)
	d.DispatchForQuota(context.Background(), event)
	if !receiver.waitFor(2, 2*time.Second) {
		t.Fatalf("re-crossing upward: webhook did not re-fire, count=%d", receiver.count())
	}
	if got := receiver.count(); got != 2 {
		t.Fatalf("re-crossing upward: expected 2 deliveries total, got %d", got)
	}
}

func TestDispatchForQuotaUnlimitedHold(t *testing.T) {
	conn, err := db.InitDB(":memory:", db.LibsqlConfig{})
	if err != nil {
		t.Fatalf("init db: %v", err)
	}
	defer conn.Close()

	const userDID = "did:plc:cap"
	if err := db.UpsertUser(conn, &db.User{DID: userDID, Handle: "cap", PDSEndpoint: "x", LastSeen: time.Now()}); err != nil {
		t.Fatal(err)
	}

	// Hold serving an unlimited response (no `limit` field)
	holdSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintf(w, `{"userDid":%q,"totalSize":9999999999}`, userDID)
	}))
	defer holdSrv.Close()

	receiver := newFakeReceiver(1)
	recvSrv := httptest.NewServer(receiver.handler())
	defer recvSrv.Close()

	if err := db.InsertWebhook(conn, &db.Webhook{
		ID: "wh", UserDID: userDID, URL: recvSrv.URL,
		Triggers: PackTriggers(TriggerQuota, 50), CreatedAt: time.Now().UTC(),
	}); err != nil {
		t.Fatal(err)
	}

	d := NewDispatcher(conn, atproto.AppviewMetadata{}, nil)
	d.allowLoopbackDeliveryForTest()
	d.DispatchForQuota(context.Background(), storage.QuotaWebhookEvent{
		UserDID: userDID, HoldDID: holdSrv.URL, HoldEndpoint: holdSrv.URL,
	})
	time.Sleep(100 * time.Millisecond)
	if got := receiver.count(); got != 0 {
		t.Fatalf("unlimited hold should never fire quota webhook, got %d deliveries", got)
	}
}