3x-ui-agent/internal/web/service/probenotify/probenotify_test.go

package probenotify

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"testing"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/web/service/probe"
)

// fakeEnv replaces every store seam of the package for one test.
type fakeEnv struct {
	settings map[string]string
	state    map[string]string
	stateErr map[string]bool
	nodes    []nodeInfo
	channels []channel
	saved    []probe.Meta
}

func newFakeEnv(t *testing.T) *fakeEnv {
	t.Helper()
	e := &fakeEnv{
		settings: map[string]string{},
		state:    map[string]string{},
		stateErr: map[string]bool{},
	}
	oldSetting, oldSetSetting := settingFn, setSettingFn
	oldChannels, oldLoad := channelsFn, loadNodesFn
	oldGet, oldSet, oldSave := stateGetFn, stateSetFn, saveMetaFn
	settingFn = func(key string) (string, bool) {
		v, ok := e.settings[key]
		return v, ok
	}
	setSettingFn = func(key, value string) error {
		if value == "" {
			delete(e.settings, key)
		} else {
			e.settings[key] = value
		}
		return nil
	}
	channelsFn = func() []channel { return e.channels }
	loadNodesFn = func() ([]nodeInfo, error) { return e.nodes, nil }
	stateGetFn = func(key string) (string, bool, error) {
		if e.stateErr[key] {
			return "", false, errors.New("store down")
		}
		v, ok := e.state[key]
		return v, ok, nil
	}
	stateSetFn = func(key, value string) error {
		if e.stateErr[key] {
			return errors.New("store down")
		}
		if value == "" {
			delete(e.state, key)
		} else {
			e.state[key] = value
		}
		return nil
	}
	saveMetaFn = func(m probe.Meta) error {
		e.saved = append(e.saved, m)
		return nil
	}
	t.Cleanup(func() {
		settingFn, setSettingFn = oldSetting, oldSetSetting
		channelsFn, loadNodesFn = oldChannels, oldLoad
		stateGetFn, stateSetFn, saveMetaFn = oldGet, oldSet, oldSave
	})
	return e
}

type fakeChannel struct {
	label string
	notes []Note
	fail  *failure
}

func (f *fakeChannel) name() string { return f.label }
func (f *fakeChannel) post(n Note) *failure {
	if f.fail != nil {
		return f.fail
	}
	f.notes = append(f.notes, n)
	return nil
}

var t0 = time.Date(2026, 10, 23, 12, 0, 0, 0, time.Local)

func offlineNode(id int64, name string, lastSeen time.Time, notify bool) nodeInfo {
	return nodeInfo{
		id:          id,
		name:        name,
		lastSeenSec: lastSeen.Unix(),
		meta:        probe.Meta{AgentID: id, Notify: notify, TrafficResetDay: 1},
	}
}

func TestSweepOfflineAfterGraceOnce(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1" // 60s
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	e.nodes = []nodeInfo{offlineNode(7, "tokyo", t0.Add(-300*time.Second), true)}
	w := newWatch()

	if notes := runSweep(t0, w); len(notes) != 0 {
		t.Fatalf("first sighting must only start the absence clock, got %v", notes)
	}
	if notes := runSweep(t0.Add(30*time.Second), w); len(notes) != 0 {
		t.Fatalf("before grace, got %v", notes)
	}
	notes := runSweep(t0.Add(61*time.Second), w)
	if len(notes) != 1 || notes[0].Event != "offline" || !strings.Contains(notes[0].Title, "tokyo") {
		t.Fatalf("after grace want one offline note, got %v", notes)
	}
	if got := e.state["notify_down:7"]; got != fmt.Sprint(t0.Add(-300*time.Second).Unix()) {
		t.Fatalf("offline mark = %q, want last-seen stamp", got)
	}
	if notes := runSweep(t0.Add(10*time.Minute), w); len(notes) != 0 {
		t.Fatalf("announced mark must mute repeats, got %v", notes)
	}
}

func TestSweepRecoveryOnceAndRearms(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1"
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	n := offlineNode(7, "tokyo", t0.Add(-300*time.Second), true)
	e.nodes = []nodeInfo{n}
	w := newWatch()
	runSweep(t0, w)
	runSweep(t0.Add(61*time.Second), w) // offline announced

	back := t0.Add(5 * time.Minute)
	n.lastSeenSec = back.Unix()
	e.nodes = []nodeInfo{n}
	notes := runSweep(back, w)
	if len(notes) != 1 || notes[0].Event != "online" || !strings.Contains(notes[0].Message, "离线") {
		t.Fatalf("want one recovery note, got %v", notes)
	}
	if _, ok := e.state["notify_down:7"]; ok {
		t.Fatal("recovery must clear the offline mark")
	}
	if notes := runSweep(back.Add(time.Minute), w); len(notes) != 0 {
		t.Fatalf("recovery must not repeat, got %v", notes)
	}
	// Down again: a fresh absence past grace announces again.
	n.lastSeenSec = back.Unix()
	e.nodes = []nodeInfo{n}
	later := back.Add(2 * time.Hour)
	if notes := runSweep(later, w); len(notes) != 0 {
		t.Fatalf("first sighting of the new absence, got %v", notes)
	}
	notes = runSweep(later.Add(61*time.Second), w)
	if len(notes) != 1 || notes[0].Event != "offline" {
		t.Fatalf("second outage must announce again, got %v", notes)
	}
}

func TestSweepFlapUsesFlapGrace(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1"
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	n := offlineNode(7, "tokyo", t0.Add(-300*time.Second), true)
	e.nodes = []nodeInfo{n}
	w := newWatch()
	runSweep(t0, w)
	runSweep(t0.Add(61*time.Second), w) // announced
	back := t0.Add(5 * time.Minute)
	n.lastSeenSec = back.Unix()
	e.nodes = []nodeInfo{n}
	runSweep(back, w) // recovered; returned[] = back

	// Away again ten minutes after returning: inside the flap
	// window, so the wait is the 30-minute flap grace, not 60s.
	away := back.Add(10 * time.Minute)
	n.lastSeenSec = back.Unix()
	e.nodes = []nodeInfo{n}
	runSweep(away, w) // first sighting of this absence
	if notes := runSweep(away.Add(61*time.Second), w); len(notes) != 0 {
		t.Fatalf("flapping absence must wait the flap grace, got %v", notes)
	}
	notes := runSweep(away.Add(flapGrace*time.Second), w)
	if len(notes) != 1 || notes[0].Event != "offline" {
		t.Fatalf("after flap grace want offline note, got %v", notes)
	}
}

func TestSweepSkipsNeverReportedAndOptedOut(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1"
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	e.nodes = []nodeInfo{
		{id: 1, name: "ghost", lastSeenSec: 0, meta: probe.Meta{AgentID: 1, Notify: true}},
		offlineNode(2, "quiet", t0.Add(-time.Hour), false),
	}
	w := newWatch()
	runSweep(t0, w)
	if notes := runSweep(t0.Add(2*time.Minute), w); len(notes) != 0 {
		t.Fatalf("never-reported and opted-out nodes stay silent, got %v", notes)
	}
	for key := range e.state {
		if strings.HasPrefix(key, "notify_down:") {
			t.Fatalf("no offline marks expected, got %v", e.state)
		}
	}
}

func TestSweepUnarmedRecordsNothing(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1"
	n := offlineNode(7, "tokyo", t0.Add(-300*time.Second), true)
	e.nodes = []nodeInfo{n}
	w := newWatch()
	runSweep(t0, w)
	if notes := runSweep(t0.Add(61*time.Second), w); len(notes) != 0 {
		t.Fatalf("no channels, no notes, got %v", notes)
	}
	if _, ok := e.state["notify_down:7"]; ok {
		t.Fatal("unarmed sweep must not record the announcement")
	}
	// Configure a channel afterwards: the already-down node is
	// reported on the next pass (its absence clock is old by now).
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	notes := runSweep(t0.Add(2*time.Minute), w)
	if len(notes) != 1 || notes[0].Event != "offline" {
		t.Fatalf("late-configured channel must still hear about the outage, got %v", notes)
	}
}

func TestSweepRestartKeepsPairing(t *testing.T) {
	e := newFakeEnv(t)
	e.settings["notify_grace"] = "1"
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	down := offlineNode(7, "tokyo", t0.Add(-300*time.Second), true)
	e.nodes = []nodeInfo{down}
	// The mark from before the "restart": announced already.
	e.state["notify_down:7"] = fmt.Sprint(down.lastSeenSec)
	w := newWatch() // fresh in-memory watch, as after a restart
	runSweep(t0, w)
	if notes := runSweep(t0.Add(2*time.Minute), w); len(notes) != 0 {
		t.Fatalf("restart must not replay the offline alert, got %v", notes)
	}
	back := t0.Add(3 * time.Minute)
	down.lastSeenSec = back.Unix()
	e.nodes = []nodeInfo{down}
	notes := runSweep(back, w)
	if len(notes) != 1 || notes[0].Event != "online" {
		t.Fatalf("the recovery pairing must survive the restart, got %v", notes)
	}
}

func trafficNode(id int64, name string, used, limit int64) nodeInfo {
	return nodeInfo{
		id:          id,
		name:        name,
		lastSeenSec: t0.Unix(),
		meta: probe.Meta{
			AgentID: id, TrafficLimit: limit, TrafficMode: "sum",
			TrafficResetDay: 1, Notify: true,
		},
		traffic:    probe.Traffic{MonthRx: used, MonthStart: "2026-10-01"},
		hasTraffic: true,
	}
}

func TestTrafficStepsPrimeThenAnnounce(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 850, 1000)}
	w := newWatch()

	// Unprimed (no stored state): the crossing is recorded silently.
	if notes := runSweep(t0, w); len(notes) != 0 {
		t.Fatalf("first pass must prime, not announce, got %v", notes)
	}
	stored := map[int64]alertStep{}
	if err := json.Unmarshal([]byte(e.state["traffic_alerted"]), &stored); err != nil {
		t.Fatalf("state not JSON: %v", err)
	}
	if stored[7].Step != 80 || stored[7].Period != "2026-10-01" {
		t.Fatalf("primed step = %+v", stored[7])
	}
	// Same step again: told already.
	if notes := runSweep(t0.Add(time.Minute), w); len(notes) != 0 {
		t.Fatalf("same step must not repeat, got %v", notes)
	}
	// Cross to 100%: the second step announces.
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 1000, 1000)}
	notes := runSweep(t0.Add(2*time.Minute), w)
	if len(notes) != 1 || notes[0].Event != "traffic" || !strings.Contains(notes[0].Message, "100%") {
		t.Fatalf("want one 100%% traffic note, got %v", notes)
	}
	if notes := runSweep(t0.Add(3*time.Minute), w); len(notes) != 0 {
		t.Fatalf("100%% step must not repeat within the period, got %v", notes)
	}
	// Next period: the same usage announces again.
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 900, 1000)}
	nov := time.Date(2026, 11, 15, 12, 0, 0, 0, time.Local)
	if notes := runSweep(nov, newWatch()); len(notes) != 1 || notes[0].Event != "traffic" {
		t.Fatalf("a new period re-arms the steps, got %v", notes)
	}
}

func TestTrafficCorrectionRearms(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	e.state["traffic_alerted"] = `{"7":["2026-10-01",80]}`
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 500, 1000)} // corrected down
	w := newWatch()
	if notes := runSweep(t0, w); len(notes) != 0 {
		t.Fatalf("a correction announces nothing, got %v", notes)
	}
	stored := map[int64]alertStep{}
	_ = json.Unmarshal([]byte(e.state["traffic_alerted"]), &stored)
	if stored[7].Step != 0 {
		t.Fatalf("the lower step must be stored, got %+v", stored[7])
	}
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 850, 1000)}
	if notes := runSweep(t0.Add(time.Minute), w); len(notes) != 1 || notes[0].Event != "traffic" {
		t.Fatalf("crossing again after a correction must announce, got %v", notes)
	}
}

func TestTrafficReadErrorSkipsSection(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	e.stateErr["traffic_alerted"] = true
	e.nodes = []nodeInfo{trafficNode(7, "tokyo", 1000, 1000)}
	if notes := runSweep(t0, newWatch()); len(notes) != 0 {
		t.Fatalf("a failed state read must skip the section, got %v", notes)
	}
}

func renewNode(id int64, name string, expires string, cycle string, online bool) nodeInfo {
	exp := expires
	ls := t0.Add(-time.Hour).Unix()
	if online {
		ls = t0.Unix()
	}
	return nodeInfo{
		id: id, name: name, lastSeenSec: ls,
		meta: probe.Meta{AgentID: id, ExpiresAt: &exp, BillingCycle: cycle, TrafficResetDay: 1},
	}
}

func TestHousekeepingRenewsOnlineNode(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	past := t0.AddDate(0, 0, -40).Format("2006-01-02")
	e.nodes = []nodeInfo{renewNode(7, "tokyo", past, "monthly", true)}
	notes := runHousekeeping(t0)
	if len(e.saved) != 1 {
		t.Fatalf("want one rolled meta, saved %v", e.saved)
	}
	want := t0.AddDate(0, 0, -40).AddDate(0, 2, 0).Format("2006-01-02")
	if *e.saved[0].ExpiresAt != want {
		t.Fatalf("rolled to %s, want %s (two whole cycles)", *e.saved[0].ExpiresAt, want)
	}
	if len(notes) != 1 || notes[0].Event != "renew" {
		t.Fatalf("want one renewal note, got %v", notes)
	}
	// Not past any more: no second roll, no second note.
	e.nodes = []nodeInfo{renewNode(7, "tokyo", *e.saved[0].ExpiresAt, "monthly", true)}
	e.saved = nil
	if notes := runHousekeeping(t0); len(notes) != 0 || len(e.saved) != 0 {
		t.Fatalf("a future date must not roll, notes %v saved %v", notes, e.saved)
	}
}

func TestHousekeepingRenewalSkipsAndGate(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	past := t0.AddDate(0, 0, -10).Format("2006-01-02")
	e.nodes = []nodeInfo{
		renewNode(1, "once-node", past, "once", true),     // one-off plan: no cycle
		renewNode(2, "away-node", past, "monthly", false), // offline: no roll
		renewNode(3, "live-node", past, "monthly", true),
	}
	e.settings["notify_expiry"] = "0" // renewals still roll, silently
	notes := runHousekeeping(t0)
	if len(e.saved) != 1 || e.saved[0].AgentID != 3 {
		t.Fatalf("only the online monthly node rolls, saved %v", e.saved)
	}
	if len(notes) != 0 {
		t.Fatalf("notify_expiry=0 must mute the renewal note, got %v", notes)
	}
}

func TestExpiryDigestOncePerDay(t *testing.T) {
	e := newFakeEnv(t)
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	soon := t0.AddDate(0, 0, 3).Format("2006-01-02")
	e.nodes = []nodeInfo{renewNode(7, "tokyo", soon, "monthly", true)}

	morning := time.Date(2026, 10, 23, 8, 0, 0, 0, time.Local)
	if notes := runHousekeeping(morning); len(notes) != 0 {
		t.Fatalf("before 09:00 no digest, got %v", notes)
	}
	if _, ok := e.state["notify_expiry_sent"]; ok {
		t.Fatal("an ungated hour must not consume the day")
	}
	notes := runHousekeeping(t0) // 12:00
	found := false
	for _, n := range notes {
		if n.Event == "expiry" && strings.Contains(n.Message, soon) {
			found = true
		}
	}
	if !found {
		t.Fatalf("want the expiry digest, got %v", notes)
	}
	if notes := runHousekeeping(t0.Add(time.Hour)); len(notes) != 0 {
		t.Fatalf("the digest is once a day, got %v", notes)
	}
}

func TestExpiryDigestNeedsChannel(t *testing.T) {
	e := newFakeEnv(t)
	soon := t0.AddDate(0, 0, 3).Format("2006-01-02")
	e.nodes = []nodeInfo{renewNode(7, "tokyo", soon, "monthly", true)}
	if notes := runHousekeeping(t0); len(notes) != 0 {
		t.Fatalf("no channel, no digest, got %v", notes)
	}
	if _, ok := e.state["notify_expiry_sent"]; ok {
		t.Fatal("the day must stay unconsumed so a channel configured later still gets the digest")
	}
	e.channels = []channel{&fakeChannel{label: "Webhook"}}
	if notes := runHousekeeping(t0); len(notes) != 1 || notes[0].Event != "expiry" {
		t.Fatalf("want the digest once a channel exists, got %v", notes)
	}
}

func TestCycleMonthsAndClamp(t *testing.T) {
	for cycle, want := range map[string]int64{
		"monthly": 1, "quarterly": 3, "semiannual": 6, "yearly": 12,
		"biennial": 24, "triennial": 36, "5m": 5, "1200m": 1200,
	} {
		if got, ok := cycleMonths(cycle); !ok || got != want {
			t.Errorf("cycleMonths(%q) = %d,%v want %d", cycle, got, ok, want)
		}
	}
	for _, bad := range []string{"once", "", "0m", "1201m", "weekly", "m"} {
		if _, ok := cycleMonths(bad); ok {
			t.Errorf("cycleMonths(%q) must not renew", bad)
		}
	}
	jan31 := time.Date(2024, 1, 31, 0, 0, 0, 0, time.Local)
	if got := addMonthsClamped(jan31, 1); got.Format("2006-01-02") != "2024-02-29" {
		t.Errorf("Jan 31 + 1 month = %v, want 2024-02-29 (clamped)", got)
	}
}

func TestSettingError(t *testing.T) {
	cases := []struct {
		key, value string
		wantErr    bool
	}{
		{"notify_grace", "3", false}, {"notify_grace", "0", true},
		{"notify_grace", "31", true}, {"notify_grace", "abc", true},
		{"notify_traffic", "0", false}, {"notify_traffic", "100", false},
		{"notify_traffic", "101", true},
		{"notify_expiry", "0", false}, {"notify_expiry", "365", false},
		{"notify_expiry", "366", true},
		{"notify_login", "on", false}, {"notify_login", "off", false},
		{"notify_login", "yes", true},
		{"notify_webhook_url", "", false},
		{"notify_webhook_url", "https://example.com/hook", false},
		{"notify_webhook_url", "http://127.0.0.1:9/x", false},
		{"notify_webhook_url", "ftp://example.com", true},
		{"notify_webhook_headers", "", false},
		{"notify_webhook_headers", "X-Api-Key: secret", false},
		{"notify_webhook_headers", "Broken Line", true},
		{"notify_webhook_headers", "Bad Name: v", true},
		{"notify_webhook_body", "", false},
		{"notify_webhook_body", `{"text":"{{title}}"}`, false},
		{"notify_webhook_body", `{"text":{{title}}}`, true},
		{"notify_webhook_body", `not json`, true},
		{"notify_telegram_text", "", false},
		{"notify_telegram_text", "anything {{title}}", false},
		{"notify_unknown", "1", true},
	}
	for _, c := range cases {
		if got := SettingError(c.key, c.value); (got != "") != c.wantErr {
			t.Errorf("SettingError(%q, %q) = %q, wantErr %v", c.key, c.value, got, c.wantErr)
		}
	}
}

func TestRenderSinglePassAndEscapes(t *testing.T) {
	note := Note{Event: "offline", Node: "a\"b", Title: "{{message}}", Message: "line\nline", Time: t0.Unix()}
	// Single pass: the title's literal {{message}} must survive.
	if got := render("{{title}}", note, "site", false); got != "{{message}}" {
		t.Errorf("render single pass = %q", got)
	}
	got := render(`{"t":"{{title}}","m":"{{message}}","n":"{{node}}"}`, note, "site", true)
	var parsed map[string]string
	if err := json.Unmarshal([]byte(got), &parsed); err != nil {
		t.Fatalf("rendered body not JSON: %v (%s)", err, got)
	}
	if parsed["m"] != "line\nline" || parsed["n"] != "a\"b" {
		t.Errorf("escapes broken: %v", parsed)
	}
	if got := render("x{{bogus}}y", note, "site", false); got != "x{{bogus}}y" {
		t.Errorf("unknown placeholder must pass through, got %q", got)
	}
}

func TestBatchShapes(t *testing.T) {
	if _, ok := batch("offline", "🔴", "离线", nil); ok {
		t.Error("empty batch must be no note")
	}
	one, _ := batch("offline", "🔴", "离线", []batchItem{{"a", "最后上报 10-23 11:00 +08:00"}})
	if one.Title != "🔴 a 离线" || one.Node != "a" {
		t.Errorf("single batch = %+v", one)
	}
	items := make([]batchItem, 25)
	for i := range items {
		items[i] = batchItem{fmt.Sprint("n", i), "d"}
	}
	many, _ := batch("offline", "🔴", "离线", items)
	if !strings.Contains(many.Title, "25 台节点离线") || !strings.Contains(many.Message, "……另外 5 台") {
		t.Errorf("many batch = %q / %q", many.Title, many.Message)
	}
}

func TestSaveSettingsValidationAndSecrets(t *testing.T) {
	e := newFakeEnv(t)
	// All-or-nothing: one bad field refuses the whole write.
	body := map[string]json.RawMessage{
		"notifyGrace":   json.RawMessage("5"),
		"notifyTraffic": json.RawMessage("101"),
	}
	if err := SaveSettings(body); err == nil {
		t.Fatal("out-of-range traffic must refuse")
	}
	if _, ok := e.settings["notify_grace"]; ok {
		t.Fatal("a refused save must write nothing")
	}
	// Secrets: null leaves, string sets, "" clears.
	e.settings["notify_webhook_url"] = "https://old.example/hook"
	if err := SaveSettings(map[string]json.RawMessage{"notifyWebhookUrl": json.RawMessage("null")}); err != nil {
		t.Fatal(err)
	}
	if e.settings["notify_webhook_url"] != "https://old.example/hook" {
		t.Fatal("null must leave the secret")
	}
	if err := SaveSettings(map[string]json.RawMessage{"notifyWebhookUrl": json.RawMessage(`""`)}); err != nil {
		t.Fatal(err)
	}
	if _, ok := e.settings["notify_webhook_url"]; ok {
		t.Fatal("empty string must clear the secret")
	}
	if err := SaveSettings(map[string]json.RawMessage{"bogus": json.RawMessage("1")}); err == nil {
		t.Fatal("unknown keys are refused")
	}
}