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, _ int) (*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, uuid.UUID, error) { if f.retryErr != nil { return 0, uuid.Nil, f.retryErr } return msgID + 100, uuid.Nil, 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") }