xinmi/internal/handler/custom_agent_draft_test.go

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