llmgw/internal/notify/notify_test.go

package notify

import (
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

type got struct {
	mu    sync.Mutex
	calls []*http.Request
	body  [][]byte
}

func (g *got) handler(fail int) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		b, _ := io.ReadAll(r.Body)
		g.mu.Lock()
		defer g.mu.Unlock()
		g.calls = append(g.calls, r)
		g.body = append(g.body, b)
		if len(g.calls) <= fail {
			w.WriteHeader(http.StatusBadGateway)
		}
	}
}

func newNotifier(t *testing.T, c Config) *Notifier {
	t.Helper()
	n, err := New(context.Background(), c, nil)
	if err != nil {
		t.Fatal(err)
	}
	n.retry = time.Millisecond
	return n
}

func TestSignedJSON(t *testing.T) {
	var g got
	srv := httptest.NewServer(g.handler(0))
	defer srv.Close()
	secret := filepath.Join(t.TempDir(), "s")
	os.WriteFile(secret, []byte("k3y\n"), 0o600)

	n := newNotifier(t, Config{Host: "gw1", Webhooks: []Webhook{{URL: srv.URL, SecretFile: secret}}})
	n.Notify(ModelParked, "qwen", "model qwen is parked", map[string]any{"failures": 3})
	n.Wait()

	if len(g.calls) != 1 {
		t.Fatalf("%d calls, want 1", len(g.calls))
	}
	m := hmac.New(sha256.New, []byte("k3y"))
	m.Write(g.body[0])
	if want := "sha256=" + hex.EncodeToString(m.Sum(nil)); g.calls[0].Header.Get("X-Llmgw-Signature") != want {
		t.Errorf("signature %q, want %q", g.calls[0].Header.Get("X-Llmgw-Signature"), want)
	}
	if g.calls[0].Header.Get("X-Llmgw-Event") != ModelParked {
		t.Errorf("event header %q", g.calls[0].Header.Get("X-Llmgw-Event"))
	}
	var e Event
	if err := json.Unmarshal(g.body[0], &e); err != nil {
		t.Fatal(err)
	}
	if e.Kind != ModelParked || e.Subject != "qwen" || e.Detail["gateway"] != "gw1" || e.Detail["failures"] != float64(3) {
		t.Errorf("event %+v", e)
	}
}

func TestSlackFiltersAndDedupes(t *testing.T) {
	var g got
	srv := httptest.NewServer(g.handler(0))
	defer srv.Close()

	n := newNotifier(t, Config{Host: "gw1", Webhooks: []Webhook{{URL: srv.URL, Format: "slack", Events: []string{NodeLost}}}})
	n.Notify(ModelParked, "qwen", "not for this hook", nil)
	n.Notify(NodeLost, "n1", "node n1 was lost", nil)
	n.Notify(NodeLost, "n1", "node n1 was lost again", nil) // within the hour: not sent
	n.Notify(NodeLost, "n2", "node n2 was lost", nil)
	n.Wait()

	if len(g.calls) != 2 {
		t.Fatalf("%d calls, want 2: %q", len(g.calls), g.body)
	}
	var m map[string]string
	json.Unmarshal(g.body[0], &m)
	if !strings.Contains(m["text"], "gw1") || !strings.Contains(m["text"], "node n1 was lost") {
		t.Errorf("text %q", m["text"])
	}
	if g.calls[0].Header.Get("X-Llmgw-Signature") != "" {
		t.Error("signed without a secret")
	}
}

func TestRetries(t *testing.T) {
	var g got
	srv := httptest.NewServer(g.handler(2))
	defer srv.Close()
	n := newNotifier(t, Config{Webhooks: []Webhook{{URL: srv.URL}}})
	n.Notify(BudgetSpent, "org:1", "spent", nil)
	n.Wait()
	if len(g.calls) != 3 {
		t.Fatalf("%d attempts, want 3 (two failures, then delivered)", len(g.calls))
	}

	var h got
	srv2 := httptest.NewServer(h.handler(10))
	defer srv2.Close()
	n = newNotifier(t, Config{Webhooks: []Webhook{{URL: srv2.URL}}})
	n.Notify(BudgetSpent, "org:1", "spent", nil)
	n.Wait()
	if len(h.calls) != 3 {
		t.Fatalf("%d attempts at a failing hook, want 3 then given up", len(h.calls))
	}
}

func TestNilAndValidate(t *testing.T) {
	var n *Notifier
	n.Notify(NodeLost, "x", "y", nil) // must not panic
	n.Wait()
	if n, err := New(context.Background(), Config{}, nil); n != nil || err != nil {
		t.Fatalf("nothing configured: %v %v", n, err)
	}
	for _, c := range []Config{
		{Webhooks: []Webhook{{URL: "ftp://x"}}},
		{Webhooks: []Webhook{{URL: "https://x", Format: "teams"}}},
		{Webhooks: []Webhook{{URL: "https://x", Events: []string{"model_gone"}}}},
	} {
		if c.Validate() == nil {
			t.Errorf("%+v accepted", c)
		}
	}
	if _, err := New(context.Background(), Config{Webhooks: []Webhook{{URL: "https://x", SecretFile: "/nonexistent"}}}, nil); err == nil {
		t.Error("a missing secret file accepted")
	}
}

func TestTestSaysHowEachAnsweredWithoutTheSecretPath(t *testing.T) {
	var g got
	ok := httptest.NewServer(g.handler(0))
	defer ok.Close()
	bad := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusForbidden) }))
	defer bad.Close()
	n := newNotifier(t, Config{Webhooks: []Webhook{
		{URL: ok.URL + "/services/T0/B0/s3cret"},
		{URL: bad.URL + "/hook"},
		{URL: "http://127.0.0.1:1/services/T0/B0/s3cret"},
	}})
	defer n.Wait()
	res := n.Test(context.Background())
	if len(res) != 3 || res[0].Err != "" || !strings.Contains(res[1].Err, "403") || res[2].Err == "" {
		t.Fatalf("%+v", res)
	}
	for _, r := range res {
		if strings.Contains(r.URL+r.Err, "s3cret") {
			t.Errorf("the secret path shown: %+v", r)
		}
	}
}

// sharedClaims is a store two control planes share.
type sharedClaims struct {
	mu   sync.Mutex
	last map[string]time.Time
	err  error
}

func (s *sharedClaims) ClaimNotice(_ context.Context, key string, at time.Time, window time.Duration) (bool, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.err != nil {
		return false, s.err
	}
	if prev, ok := s.last[key]; ok && at.Sub(prev) < window {
		return false, nil
	}
	s.last[key] = at
	return true, nil
}

func TestTwoControlPlanesSendAnEventOnce(t *testing.T) {
	// Each control plane serving a team's requests sees its budget spent,
	// and each kept its own hour of quiet: two behind a balancer sent every
	// budget event twice.
	var g got
	srv := httptest.NewServer(g.handler(0))
	defer srv.Close()
	claims := &sharedClaims{last: map[string]time.Time{}}
	a := newNotifier(t, Config{Host: "ctl1", Webhooks: []Webhook{{URL: srv.URL}}})
	b := newNotifier(t, Config{Host: "ctl2", Webhooks: []Webhook{{URL: srv.URL}}})
	a.SetClaimer(claims)
	b.SetClaimer(claims)
	a.Notify(BudgetSpent, "team:t1", "spent", nil)
	b.Notify(BudgetSpent, "team:t1", "spent", nil)
	b.Notify(BudgetSpent, "team:t2", "spent", nil)
	a.Wait()
	b.Wait()
	if len(g.calls) != 2 {
		t.Errorf("%d sent, want one for each team", len(g.calls))
	}

	// A store that cannot be asked: sent, rather than lost.
	var g2 got
	srv2 := httptest.NewServer(g2.handler(0))
	defer srv2.Close()
	c := newNotifier(t, Config{Webhooks: []Webhook{{URL: srv2.URL}}})
	c.SetClaimer(&sharedClaims{err: errors.New("store down")})
	c.Notify(NodeLost, "node3", "lost", nil)
	c.Wait()
	if len(g2.calls) != 1 {
		t.Errorf("with the store down, %d sent, want 1", len(g2.calls))
	}
}