51 lines
1.7 KiB
Go
51 lines
1.7 KiB
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)
|
||
|
|
}
|
||
|
|
}
|