xinmi/internal/workbench/operations_diagnostics_test.go

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)
		}
	}
}