package workbench_test import ( "context" "encoding/json" "errors" "fmt" "net/http" "net/http/httptest" "path/filepath" "strings" "testing" apperrors "github.com/Tencent/WeKnora/internal/errors" "github.com/Tencent/WeKnora/internal/middleware" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/workbench" "github.com/gin-gonic/gin" "gorm.io/driver/sqlite" "gorm.io/gorm" "gorm.io/gorm/logger" ) func operationTestDB(t *testing.T) *gorm.DB { t.Helper() db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "operations.db")), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)}) if err != nil { t.Fatal(err) } sqlDB, err := db.DB() if err != nil { t.Fatal(err) } t.Cleanup(func() { _ = sqlDB.Close() }) if err := db.AutoMigrate(&workbench.OperationLog{}); err != nil { t.Fatal(err) } return db } func operationTestContext() context.Context { ctx := context.WithValue(context.Background(), types.TenantIDContextKey, uint64(1)) ctx = context.WithValue(ctx, types.UserIDContextKey, "alice") return context.WithValue(ctx, types.TenantRoleContextKey, types.TenantRoleAdmin) } func TestOperationAuditUsesRenderedErrorStatus(t *testing.T) { gin.SetMode(gin.TestMode) for _, tc := range []struct { name, method, level string err error written, status int rows int }{ {"read forbidden", "GET", "warn", fmt.Errorf("wrapped: %w", apperrors.NewForbiddenError("无权访问")), 0, 403, 2}, {"write failed", "POST", "error", errors.New("private prompt and secret-api-key"), 0, 500, 2}, {"read failed", "GET", "error", errors.New("private prompt"), 0, 500, 2}, {"written response", "POST", "info", errors.New("late error"), 202, 202, 1}, {"successful read", "GET", "", nil, 200, 200, 0}, } { t.Run(tc.name, func(t *testing.T) { db := operationTestDB(t) h := workbench.NewHandler(workbench.NewService(db)) router := gin.New() router.Use(middleware.ErrorHandler(), h.OperationAudit()) router.Handle(tc.method, "/agents/:id", func(c *gin.Context) { if tc.err != nil { _ = c.Error(tc.err) } if tc.written > 0 { c.JSON(tc.written, gin.H{"retained": true}) } }) request := httptest.NewRequest(tc.method, "/agents/agent-1", nil).WithContext(operationTestContext()) response := httptest.NewRecorder() router.ServeHTTP(response, request) if response.Code != tc.status || !json.Valid(response.Body.Bytes()) { t.Fatalf("invalid rendered response: %d %s", response.Code, response.Body.String()) } var rows []workbench.OperationLog if err := db.Find(&rows).Error; err != nil { t.Fatal(err) } if len(rows) != tc.rows { t.Fatalf("unexpected log count: %+v", rows) } for _, row := range rows { if row.Status != tc.status || row.Level != tc.level || row.TenantID != 1 || row.OwnerID != "alice" || row.ResourceID != "agent-1" { t.Fatalf("incorrect final status or ownership: %+v", row) } serialized, _ := json.Marshal(row) if strings.Contains(string(serialized), "private prompt") || strings.Contains(string(serialized), "secret-api-key") { t.Fatal("raw error leaked into operation log") } } }) } } func TestSystemLogQueryHidesWarningsAndOtherTenants(t *testing.T) { db := operationTestDB(t) for _, row := range []workbench.OperationLog{ {ID: "error", TenantID: 1, Kind: "system", Level: "error"}, {ID: "warn", TenantID: 1, Kind: "system", Level: "warn"}, {ID: "warning", TenantID: 1, Kind: "system", Level: "warning"}, {ID: "other-tenant", TenantID: 2, Kind: "system", Level: "error"}, } { if err := db.Create(&row).Error; err != nil { t.Fatal(err) } } router := gin.New() router.GET("/operations/:recordKind", workbench.NewHandler(workbench.NewService(db)).OperationRecords) request := httptest.NewRequest(http.MethodGet, "/operations/system", nil).WithContext(operationTestContext()) response := httptest.NewRecorder() router.ServeHTTP(response, request) var body struct { Data []workbench.OperationLog `json:"data"` Total int `json:"total"` } if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { t.Fatal(err) } if response.Code != 200 || body.Total != 1 || len(body.Data) != 1 || body.Data[0].ID != "error" { t.Fatalf("system log filtering failed: %s", response.Body.String()) } } func TestWorkbenchErrorRetainsResponseAndRecordsDiagnostic(t *testing.T) { db := operationTestDB(t) h := workbench.NewHandler(workbench.NewService(db)) router := gin.New() router.Use(middleware.ErrorHandler(), h.OperationAudit()) router.GET("/workbench/:kind", h.List) request := httptest.NewRequest(http.MethodGet, "/workbench/unknown-kind", nil).WithContext(operationTestContext()) response := httptest.NewRecorder() router.ServeHTTP(response, request) var body map[string]any if err := json.Unmarshal(response.Body.Bytes(), &body); err != nil { t.Fatalf("response was rendered more than once or is invalid: %s", response.Body.String()) } if response.Code != http.StatusBadRequest || body["success"] != false || body["message"] != workbench.ErrInvalid.Error() || len(body) != 2 { t.Fatalf("existing error response changed: %d %s", response.Code, response.Body.String()) } var rows []workbench.OperationLog if err := db.Find(&rows).Error; err != nil { t.Fatal(err) } if len(rows) != 2 { t.Fatalf("expected audit and system records, got %+v", rows) } for _, row := range rows { if row.Status != http.StatusBadRequest || row.Metadata["error_class"] != "invalid" || row.Metadata["error_message"] != workbench.ErrInvalid.Error() || row.Level != "warn" { t.Fatalf("underlying error was not attached to logs: %+v", row) } } }