Add files via upload

internal/handler/vulnerability_alert_test.go

@@ -0,0 +1,51 @@
+package handler
+
+import (
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"cyberstrike-ai/internal/config"
+	"cyberstrike-ai/internal/database"
+	"cyberstrike-ai/internal/security"
+
+	"go.uber.org/zap"
+)
+
+func TestRobotVulnerabilityAlertCommandSharesSubscription(t *testing.T) {
+	db, err := database.NewDB(filepath.Join(t.TempDir(), "alert-command.db"), zap.NewNop())
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = db.Close() })
+	if err := db.BootstrapRBAC("hash", security.PermissionCatalog); err != nil {
+		t.Fatal(err)
+	}
+	user, err := db.CreateRBACUser("alert-user", "Alert User", "hash", true, []string{database.RBACSystemRoleOperator})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if err := db.CreateRobotBindingCode(user.ID, "alert-code", time.Now().Add(time.Minute)); err != nil {
+		t.Fatal(err)
+	}
+	if _, err := db.ConsumeRobotBindingCode("wecom", "external-user", "alert-code"); err != nil {
+		t.Fatal(err)
+	}
+	h := NewRobotHandler(&config.Config{}, db, nil, zap.NewNop())
+
+	if got := h.HandleMessage("wecom", "external-user", "漏洞提醒 高危以上"); !strings.Contains(got, "已开启") {
+		t.Fatalf("enable reply: %s", got)
+	}
+	sub, err := db.GetVulnerabilityAlertSubscription(user.ID)
+	if err != nil || !sub.Enabled || sub.MinSeverity != "high" {
+		t.Fatalf("web subscription not updated: %#v %v", sub, err)
+	}
+	if got := h.HandleMessage("wecom", "external-user", "vuln alerts off"); !strings.Contains(got, "已关闭") {
+		t.Fatalf("disable reply: %s", got)
+	}
+	sub, _ = db.GetVulnerabilityAlertSubscription(user.ID)
+	if sub.Enabled {
+		t.Fatalf("subscription remained enabled: %#v", sub)
+	}
+}