package handler import ( "context" "net/http" "net/http/httptest" "strings" "testing" "github.com/Tencent/WeKnora/internal/types" "github.com/Tencent/WeKnora/internal/types/interfaces" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) type draftHandlerService struct { interfaces.CustomAgentService calls int } func (s *draftHandlerService) DraftAgent(ctx context.Context, model, requirement string) (*types.CustomAgent, error) { s.calls++ return &types.CustomAgent{Name: requirement, Config: types.CustomAgentConfig{ModelID: model}}, nil } func TestAgentDraftHandlerReturnsUnsavedEditableDraft(t *testing.T) { svc := &draftHandlerService{} h := &CustomAgentHandler{service: svc} w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/agents/draft", strings.NewReader(`{"model_id":"selected","requirement":"Plan a course"}`)) c.Request.Header.Set("Content-Type", "application/json") h.DraftAgent(c) require.Equal(t, http.StatusOK, w.Code) require.Equal(t, 1, svc.calls) require.Contains(t, w.Body.String(), `"model_id":"selected"`) require.Contains(t, w.Body.String(), `"id":""`) } func TestAgentDraftHandlerRejectsMissingAndOversizedInput(t *testing.T) { for _, body := range []string{`{"requirement":"Plan"}`, `{"model_id":"selected","requirement":"` + strings.Repeat("x", 33000) + `"}`} { svc := &draftHandlerService{} c, _ := gin.CreateTestContext(httptest.NewRecorder()) c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/agents/draft", strings.NewReader(body)) c.Request.Header.Set("Content-Type", "application/json") (&CustomAgentHandler{service: svc}).DraftAgent(c) require.NotEmpty(t, c.Errors) require.Zero(t, svc.calls) } }