You've already forked agentic-coding-workflow
391 lines
14 KiB
Go
391 lines
14 KiB
Go
package chat
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/go-chi/chi/v5"
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/assert"
|
|
"github.com/stretchr/testify/require"
|
|
|
|
"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
|
|
"github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware"
|
|
)
|
|
|
|
// ===== Fake service implementations for handler tests =====
|
|
|
|
// fakeSessionResolver always resolves to fakeUID when a non-empty token is provided.
|
|
type fakeSessionResolver struct {
|
|
uid uuid.UUID
|
|
}
|
|
|
|
func (f *fakeSessionResolver) ResolveSession(_ context.Context, token string) (uuid.UUID, bool, error) {
|
|
if token == "" {
|
|
return uuid.Nil, false, nil
|
|
}
|
|
return f.uid, true, nil
|
|
}
|
|
|
|
// fakeAdminLookup allows controlling IsAdmin per-user.
|
|
type fakeAdminLookup struct {
|
|
adminIDs map[uuid.UUID]bool
|
|
}
|
|
|
|
func newFakeAdminLookup(admins ...uuid.UUID) *fakeAdminLookup {
|
|
m := make(map[uuid.UUID]bool, len(admins))
|
|
for _, id := range admins {
|
|
m[id] = true
|
|
}
|
|
return &fakeAdminLookup{adminIDs: m}
|
|
}
|
|
|
|
func (f *fakeAdminLookup) IsAdmin(_ context.Context, id uuid.UUID) (bool, error) {
|
|
return f.adminIDs[id], nil
|
|
}
|
|
|
|
// ===== Fake ConversationService =====
|
|
|
|
type fakeConvSvc struct {
|
|
createFunc func(ctx context.Context, userID uuid.UUID, in CreateConversationInput) (*Conversation, error)
|
|
listFunc func(ctx context.Context, userID uuid.UUID, f ConversationFilter) ([]Conversation, error)
|
|
getFunc func(ctx context.Context, userID, id uuid.UUID) (*Conversation, []Message, error)
|
|
}
|
|
|
|
func (f *fakeConvSvc) Create(ctx context.Context, userID uuid.UUID, in CreateConversationInput) (*Conversation, error) {
|
|
return f.createFunc(ctx, userID, in)
|
|
}
|
|
func (f *fakeConvSvc) List(ctx context.Context, userID uuid.UUID, filter ConversationFilter) ([]Conversation, error) {
|
|
if f.listFunc != nil {
|
|
return f.listFunc(ctx, userID, filter)
|
|
}
|
|
return nil, nil
|
|
}
|
|
func (f *fakeConvSvc) Get(ctx context.Context, userID, id uuid.UUID) (*Conversation, []Message, error) {
|
|
if f.getFunc != nil {
|
|
return f.getFunc(ctx, userID, id)
|
|
}
|
|
return nil, nil, errs.New(errs.CodeChatConversationNotFound, "not found")
|
|
}
|
|
func (f *fakeConvSvc) UpdateTitle(_ context.Context, _, _ uuid.UUID, _ string) error { return nil }
|
|
func (f *fakeConvSvc) SoftDelete(_ context.Context, _, _ uuid.UUID) error { return nil }
|
|
|
|
// ===== Fake MessageService =====
|
|
|
|
type fakeMsgSvc struct {
|
|
sendFunc func(ctx context.Context, userID, convID uuid.UUID, content string, ids []uuid.UUID) (int64, int64, error)
|
|
retryErr error
|
|
}
|
|
|
|
func (f *fakeMsgSvc) Send(ctx context.Context, userID, convID uuid.UUID, content string, ids []uuid.UUID) (int64, int64, error) {
|
|
if f.sendFunc != nil {
|
|
return f.sendFunc(ctx, userID, convID, content, ids)
|
|
}
|
|
return 1, 2, nil
|
|
}
|
|
func (f *fakeMsgSvc) Stream(_ context.Context, _ uuid.UUID, _ uuid.UUID, _ int64) (<-chan SSEEvent, error) {
|
|
ch := make(chan SSEEvent)
|
|
close(ch)
|
|
return ch, nil
|
|
}
|
|
func (f *fakeMsgSvc) Cancel(_ context.Context, _ uuid.UUID, _ int64) error { return nil }
|
|
func (f *fakeMsgSvc) Retry(_ context.Context, _ uuid.UUID, msgID int64) (int64, error) {
|
|
if f.retryErr != nil {
|
|
return 0, f.retryErr
|
|
}
|
|
return msgID + 100, nil
|
|
}
|
|
func (f *fakeMsgSvc) ListMessages(_ context.Context, _, _ uuid.UUID, _ *int64, _ int) ([]Message, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
// ===== Fake TemplateService =====
|
|
|
|
type fakeTplSvc struct {
|
|
createFunc func(ctx context.Context, userID uuid.UUID, name, content string, scope TemplateScope) (*PromptTemplate, error)
|
|
}
|
|
|
|
func (f *fakeTplSvc) List(_ context.Context, _ uuid.UUID) ([]PromptTemplate, error) { return nil, nil }
|
|
func (f *fakeTplSvc) Create(ctx context.Context, userID uuid.UUID, name, content string, scope TemplateScope) (*PromptTemplate, error) {
|
|
if f.createFunc != nil {
|
|
return f.createFunc(ctx, userID, name, content, scope)
|
|
}
|
|
return &PromptTemplate{ID: uuid.New(), Name: name, Content: content, Scope: scope, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
|
|
}
|
|
func (f *fakeTplSvc) Update(_ context.Context, _, _ uuid.UUID, _, _ string) error { return nil }
|
|
func (f *fakeTplSvc) Delete(_ context.Context, _, _ uuid.UUID) error { return nil }
|
|
|
|
// ===== Fake EndpointService =====
|
|
|
|
type fakeEpSvc struct {
|
|
listEndpointsErr error
|
|
}
|
|
|
|
func (f *fakeEpSvc) ListEndpoints(_ context.Context) ([]LLMEndpoint, error) {
|
|
if f.listEndpointsErr != nil {
|
|
return nil, f.listEndpointsErr
|
|
}
|
|
return nil, nil
|
|
}
|
|
func (f *fakeEpSvc) CreateEndpoint(_ context.Context, _ uuid.UUID, in CreateEndpointInput) (*LLMEndpoint, error) {
|
|
return &LLMEndpoint{ID: uuid.New(), Provider: in.Provider, DisplayName: in.DisplayName, BaseURL: in.BaseURL, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
|
|
}
|
|
func (f *fakeEpSvc) UpdateEndpoint(_ context.Context, _ uuid.UUID, _ UpdateEndpointInput) error {
|
|
return nil
|
|
}
|
|
func (f *fakeEpSvc) DeleteEndpoint(_ context.Context, _ uuid.UUID) error { return nil }
|
|
func (f *fakeEpSvc) TestEndpoint(_ context.Context, _ uuid.UUID) error { return nil }
|
|
func (f *fakeEpSvc) ListModels(_ context.Context, _ bool) ([]LLMModel, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeEpSvc) CreateModel(_ context.Context, in CreateModelInput) (*LLMModel, error) {
|
|
return &LLMModel{ID: uuid.New(), ModelID: in.ModelID, DisplayName: in.DisplayName, CreatedAt: time.Now(), UpdatedAt: time.Now()}, nil
|
|
}
|
|
func (f *fakeEpSvc) UpdateModel(_ context.Context, _ uuid.UUID, _ UpdateModelInput) error {
|
|
return nil
|
|
}
|
|
func (f *fakeEpSvc) DeleteModel(_ context.Context, _ uuid.UUID) error { return nil }
|
|
|
|
// ===== Fake UsageService =====
|
|
|
|
type fakeUsageSvc struct{}
|
|
|
|
func (f *fakeUsageSvc) Insert(_ context.Context, _ UsageRecord) error { return nil }
|
|
func (f *fakeUsageSvc) SummarizeByUser(_ context.Context, _, _ time.Time) ([]UserUsageRow, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeUsageSvc) SummarizeByModel(_ context.Context, _, _ time.Time) ([]ModelUsageRow, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeUsageSvc) SummarizeByDay(_ context.Context, _, _ time.Time) ([]DayUsageRow, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeUsageSvc) UserDaily(_ context.Context, _ uuid.UUID, _, _ time.Time) ([]DayUsageRow, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
// ===== Test helper =====
|
|
|
|
// newTestHandler builds a Handler with the given services and mounts it on a chi router.
|
|
// The resolver resolves any non-empty token to uid; adminIDs are treated as admins.
|
|
func newTestHandler(
|
|
convSvc ConversationService,
|
|
msgSvc MessageService,
|
|
tplSvc TemplateService,
|
|
epSvc EndpointService,
|
|
uid uuid.UUID,
|
|
adminLookup AdminLookup,
|
|
) (*Handler, chi.Router) {
|
|
resolver := &fakeSessionResolver{uid: uid}
|
|
upload := NewUploader(newFakeRepo(), nil, nil, 1<<20)
|
|
h := NewHandler(convSvc, msgSvc, tplSvc, epSvc, &fakeUsageSvc{}, upload, resolver, adminLookup)
|
|
r := chi.NewRouter()
|
|
r.Use(middleware.RequestID)
|
|
h.Mount(r)
|
|
return h, r
|
|
}
|
|
|
|
// authRequest creates a request with a Bearer token so Auth middleware passes.
|
|
func authRequest(method, path string, body []byte) *http.Request {
|
|
var r *http.Request
|
|
if body != nil {
|
|
r = httptest.NewRequest(method, path, bytes.NewReader(body))
|
|
r.Header.Set("Content-Type", "application/json")
|
|
} else {
|
|
r = httptest.NewRequest(method, path, nil)
|
|
}
|
|
r.Header.Set("Authorization", "Bearer test-token")
|
|
return r
|
|
}
|
|
|
|
// ===== Tests =====
|
|
|
|
// TestHandler_SendMessage_Returns200WithStreamURL verifies that sendMessage
|
|
// calls the service with the right arguments and returns a SendMessageResp
|
|
// with a correctly formed StreamURL.
|
|
func TestHandler_SendMessage_Returns200WithStreamURL(t *testing.T) {
|
|
uid := uuid.New()
|
|
convID := uuid.New()
|
|
|
|
var capturedConvID uuid.UUID
|
|
var capturedContent string
|
|
msgSvc := &fakeMsgSvc{
|
|
sendFunc: func(_ context.Context, _ uuid.UUID, cID uuid.UUID, content string, _ []uuid.UUID) (int64, int64, error) {
|
|
capturedConvID = cID
|
|
capturedContent = content
|
|
return 10, 20, nil
|
|
},
|
|
}
|
|
|
|
_, r := newTestHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
body, _ := json.Marshal(SendMessageReq{Content: "hello"})
|
|
req := authRequest(http.MethodPost, "/api/v1/conversations/"+convID.String()+"/messages", body)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
require.Equal(t, http.StatusOK, w.Code)
|
|
var resp SendMessageResp
|
|
require.NoError(t, json.NewDecoder(w.Body).Decode(&resp))
|
|
assert.Equal(t, int64(10), resp.UserMessageID)
|
|
assert.Equal(t, int64(20), resp.AssistantMessageID)
|
|
assert.Contains(t, resp.StreamURL, convID.String())
|
|
assert.Equal(t, convID, capturedConvID)
|
|
assert.Equal(t, "hello", capturedContent)
|
|
}
|
|
|
|
// TestHandler_CreateConversation_LocksTemplateContent verifies that the handler
|
|
// passes the TemplateID and SystemPrompt fields through to the service unchanged.
|
|
// The service is responsible for resolving template content; the handler must
|
|
// not override or ignore these fields.
|
|
func TestHandler_CreateConversation_LocksTemplateContent(t *testing.T) {
|
|
uid := uuid.New()
|
|
templateID := uuid.New()
|
|
customPrompt := "my custom system prompt"
|
|
|
|
var capturedInput CreateConversationInput
|
|
convSvc := &fakeConvSvc{
|
|
createFunc: func(_ context.Context, _ uuid.UUID, in CreateConversationInput) (*Conversation, error) {
|
|
capturedInput = in
|
|
return &Conversation{
|
|
ID: uuid.New(),
|
|
UserID: uid,
|
|
ModelID: in.ModelID,
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
}, nil
|
|
},
|
|
}
|
|
|
|
_, r := newTestHandler(convSvc, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
modelID := uuid.New()
|
|
body, _ := json.Marshal(CreateConversationReq{
|
|
ModelID: modelID,
|
|
TemplateID: &templateID,
|
|
SystemPrompt: &customPrompt,
|
|
})
|
|
req := authRequest(http.MethodPost, "/api/v1/conversations", body)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
require.Equal(t, http.StatusCreated, w.Code)
|
|
// Handler must forward both TemplateID and SystemPrompt to the service
|
|
// (service decides precedence, not handler).
|
|
assert.Equal(t, modelID, capturedInput.ModelID)
|
|
require.NotNil(t, capturedInput.TemplateID)
|
|
assert.Equal(t, templateID, *capturedInput.TemplateID)
|
|
require.NotNil(t, capturedInput.SystemPrompt)
|
|
assert.Equal(t, customPrompt, *capturedInput.SystemPrompt)
|
|
}
|
|
|
|
// TestHandler_AdminEndpoint_403WhenNotAdmin verifies that a non-admin user
|
|
// receives 403 when calling an admin-only endpoint.
|
|
func TestHandler_AdminEndpoint_403WhenNotAdmin(t *testing.T) {
|
|
uid := uuid.New()
|
|
// uid is NOT in the admin lookup.
|
|
_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
req := authRequest(http.MethodGet, "/api/v1/admin/llm-endpoints", nil)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusForbidden, w.Code)
|
|
var body map[string]any
|
|
require.NoError(t, json.NewDecoder(w.Body).Decode(&body))
|
|
assert.Equal(t, string(errs.CodeForbidden), body["code"])
|
|
}
|
|
|
|
// TestHandler_RetryMessage_404OnUnknownID verifies that when the message
|
|
// service returns a not-found error, the handler maps it to 404.
|
|
func TestHandler_RetryMessage_404OnUnknownID(t *testing.T) {
|
|
uid := uuid.New()
|
|
msgSvc := &fakeMsgSvc{
|
|
retryErr: errs.New(errs.CodeChatMessageNotFound, "message not found"),
|
|
}
|
|
|
|
_, r := newTestHandler(&fakeConvSvc{}, msgSvc, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
req := authRequest(http.MethodPost, "/api/v1/messages/99999/retry", nil)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusNotFound, w.Code)
|
|
var body map[string]any
|
|
require.NoError(t, json.NewDecoder(w.Body).Decode(&body))
|
|
assert.Equal(t, string(errs.CodeChatMessageNotFound), body["code"])
|
|
}
|
|
|
|
// TestHandler_CreateTemplate_SystemScopeRequiresAdmin verifies that a
|
|
// non-admin user cannot create a system-scoped template and receives 403.
|
|
func TestHandler_CreateTemplate_SystemScopeRequiresAdmin(t *testing.T) {
|
|
uid := uuid.New()
|
|
// uid is not an admin.
|
|
_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
body, _ := json.Marshal(map[string]string{
|
|
"name": "sys tpl",
|
|
"content": "You are a system assistant",
|
|
"scope": "system",
|
|
})
|
|
req := authRequest(http.MethodPost, "/api/v1/prompt-templates", body)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusForbidden, w.Code)
|
|
}
|
|
|
|
// TestHandler_CreateTemplate_SystemScopeAllowedForAdmin verifies that an
|
|
// admin user can create a system-scoped template.
|
|
func TestHandler_CreateTemplate_SystemScopeAllowedForAdmin(t *testing.T) {
|
|
uid := uuid.New()
|
|
// uid is an admin.
|
|
_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup(uid))
|
|
|
|
body, _ := json.Marshal(map[string]string{
|
|
"name": "sys tpl",
|
|
"content": "You are a system assistant",
|
|
"scope": "system",
|
|
})
|
|
req := authRequest(http.MethodPost, "/api/v1/prompt-templates", body)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusCreated, w.Code)
|
|
}
|
|
|
|
// TestHandler_Unauthenticated_Returns401 verifies that requests without a
|
|
// Bearer token are rejected by the Auth middleware with 401.
|
|
func TestHandler_Unauthenticated_Returns401(t *testing.T) {
|
|
uid := uuid.New()
|
|
_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup())
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/api/v1/conversations", nil)
|
|
// No Authorization header.
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusUnauthorized, w.Code)
|
|
}
|
|
|
|
// TestHandler_AdminAllowed_ListEndpoints verifies that an admin can call a
|
|
// admin-only endpoint and receives a 200.
|
|
func TestHandler_AdminAllowed_ListEndpoints(t *testing.T) {
|
|
uid := uuid.New()
|
|
_, r := newTestHandler(&fakeConvSvc{}, &fakeMsgSvc{}, &fakeTplSvc{}, &fakeEpSvc{}, uid, newFakeAdminLookup(uid))
|
|
|
|
req := authRequest(http.MethodGet, "/api/v1/admin/llm-endpoints", nil)
|
|
w := httptest.NewRecorder()
|
|
r.ServeHTTP(w, req)
|
|
|
|
assert.Equal(t, http.StatusOK, w.Code)
|
|
// Response should be a JSON array (empty is fine).
|
|
assert.True(t, strings.HasPrefix(w.Body.String(), "[") || w.Body.String() == "[]\n")
|
|
}
|