package media
import (
"bytes"
"context"
"encoding/base64"
"errors"
"image"
"image/color"
"image/png"
"mime/multipart"
"net/http"
"net/http/httptest"
"net/textproto"
"strconv"
"testing"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
"github.com/tidbcloud/AgentTeam/internal/api"
"github.com/tidbcloud/AgentTeam/internal/store/agent"
"github.com/tidbcloud/AgentTeam/internal/store/attachment"
"github.com/tidbcloud/AgentTeam/internal/store/identity"
"github.com/tidbcloud/AgentTeam/internal/store/storeerr"
"github.com/tidbcloud/AgentTeam/internal/store/workspace"
)
// --- Fake attachment stores ---
type fakeAttachmentStore struct {
uploadResult []attachment.AppAttachment
uploadErr error
avatarID string
avatarErr error
avatarOwner string
avatarKind string
avatarWsp string
}
type fakeAgentAvatarProfileStore struct {
getErr error
updateErr error
getCalls int
updateCalls int
updatedURL string
updatedAgent string
}
func (f *fakeAgentAvatarProfileStore) GetAgentForUser(_ context.Context, _ string, _ string, agentID string) (agent.AgentView, error) {
f.getCalls++
if f.getErr != nil {
return agent.AgentView{}, f.getErr
}
return agent.AgentView{AgentID: agentID}, nil
}
func (f *fakeAgentAvatarProfileStore) UpdateAgentForUser(_ context.Context, _ string, _ string, agentID string, params agent.UpdateAgentParams) (agent.AgentView, error) {
f.updateCalls++
f.updatedAgent = agentID
if params.AvatarURL != nil {
f.updatedURL = *params.AvatarURL
}
if f.updateErr != nil {
return agent.AgentView{}, f.updateErr
}
return agent.AgentView{AgentID: agentID}, nil
}
func (f *fakeAttachmentStore) AuthenticateSession(_ context.Context, _ string) (identity.AuthenticatedUser, error) {
return identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}}, nil
}
func (f *fakeAttachmentStore) GetWorkspaceForUser(_ context.Context, _ string, _ string) (workspace.Summary, error) {
return workspace.Summary{WorkspaceID: "wsp_1"}, nil
}
type fakeWorkspaceResolver struct {
workspace workspace.Summary
err error
calls int
}
func (f *fakeWorkspaceResolver) GetWorkspaceForUser(_ context.Context, _ string, _ string) (workspace.Summary, error) {
f.calls++
if f.err != nil {
return workspace.Summary{}, f.err
}
if f.workspace.WorkspaceID == "" {
return workspace.Summary{WorkspaceID: "wsp_1"}, nil
}
return f.workspace, nil
}
func (f *fakeAttachmentStore) MemberIDForAgentAuthToken(_ context.Context, _ string) (string, error) {
return "member_1", nil
}
func (f *fakeAttachmentStore) UploadAttachmentsForUser(_ context.Context, _ string, _ string, _ string, _ []attachment.CreateAttachmentParams) ([]attachment.AppAttachment, error) {
if f.uploadErr != nil {
return nil, f.uploadErr
}
return f.uploadResult, nil
}
func (f *fakeAttachmentStore) UploadUserAvatar(_ context.Context, _ string, _ string, _ attachment.CreateAttachmentParams) (string, error) {
if f.avatarErr != nil {
return "", f.avatarErr
}
return f.avatarID, nil
}
func (f *fakeAttachmentStore) UploadAgentAvatar(_ context.Context, workspaceID string, agentID string, _ attachment.CreateAttachmentParams) (string, error) {
if f.avatarErr != nil {
return "", f.avatarErr
}
f.avatarWsp = workspaceID
f.avatarKind = "agent"
f.avatarOwner = agentID
return f.avatarID, nil
}
func (f *fakeAttachmentStore) UploadResourceAvatar(_ context.Context, workspaceID string, ownerKind string, ownerID string, _ attachment.CreateAttachmentParams) (string, error) {
if f.avatarErr != nil {
return "", f.avatarErr
}
f.avatarWsp = workspaceID
f.avatarKind = ownerKind
f.avatarOwner = ownerID
return f.avatarID, nil
}
func (f *fakeAttachmentStore) LoadUserAvatar(_ context.Context, _ string) (attachment.AttachmentDownload, error) {
return attachment.AttachmentDownload{Filename: "avatar.png", ContentType: "image/png", Blob: []byte("png-data")}, nil
}
func (f *fakeAttachmentStore) LoadAgentAvatar(_ context.Context, _ string, _ string) (attachment.AttachmentDownload, error) {
return attachment.AttachmentDownload{Filename: "agent-avatar.png", ContentType: "image/png", Blob: []byte("agent-png-data")}, nil
}
func (f *fakeAttachmentStore) LoadAvatar(_ context.Context, _ string) (attachment.AttachmentDownload, error) {
return attachment.AttachmentDownload{Filename: "avatar.png", ContentType: "image/png", Blob: []byte("png-data")}, nil
}
func (f *fakeAttachmentStore) LoadAttachmentForMember(_ context.Context, _ string, _ string, _ string) (attachment.AttachmentDownload, error) {
return attachment.AttachmentDownload{Filename: "file.txt", ContentType: "text/plain", Blob: []byte("data")}, nil
}
func (f *fakeAttachmentStore) LoadAttachmentForUser(_ context.Context, _ string, _ string, _ string) (attachment.AttachmentDownload, error) {
return attachment.AttachmentDownload{Filename: "file.txt", ContentType: "text/plain", Blob: []byte("data")}, nil
}
func (f *fakeAttachmentStore) LoadAttachmentThumbnail(_ context.Context, _ string, _ string) (attachment.AttachmentVariant, error) {
return attachment.AttachmentVariant{}, attachment.ErrNotFound
}
func (f *fakeAttachmentStore) AuthorizeAttachmentRead(_ context.Context, _ string, _ string, _ string, _ bool, _ string) error {
return nil
}
func (f *fakeAttachmentStore) SaveAttachmentThumbnail(_ context.Context, _ string, _ string, _ attachment.GeneratedThumbnail) (attachment.AttachmentVariant, error) {
return attachment.AttachmentVariant{}, attachment.ErrNotFound
}
func (f *fakeAttachmentStore) ResolveUserAvatarRef(_ context.Context, _ string) (attachment.AvatarVariantRef, error) {
return attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"}, nil
}
func (f *fakeAttachmentStore) ResolveAgentAvatarRef(_ context.Context, _ string, _ string) (attachment.AvatarVariantRef, error) {
return attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"}, nil
}
func (f *fakeAttachmentStore) LoadAvatarThumbnail(_ context.Context, _ string, _ string) (attachment.AttachmentVariant, error) {
return attachment.AttachmentVariant{}, attachment.ErrNotFound
}
func (f *fakeAttachmentStore) SaveAvatarThumbnail(_ context.Context, _ string, variant string, gen attachment.GeneratedAvatarThumbnail) (attachment.AttachmentVariant, error) {
return attachment.AttachmentVariant{ContentType: gen.ContentType, Blob: gen.Blob, Width: gen.Width, Height: gen.Height}, nil
}
type fakeAttachmentUploadErrorStore struct {
fakeAttachmentStore
}
func (f *fakeAttachmentUploadErrorStore) UploadAttachmentsForUser(_ context.Context, _ string, _ string, _ string, _ []attachment.CreateAttachmentParams) ([]attachment.AppAttachment, error) {
return nil, storeerr.ErrForbidden
}
type fakeAvatarUploadErrorStore struct {
fakeAttachmentStore
}
func (f *fakeAvatarUploadErrorStore) UploadUserAvatar(_ context.Context, _ string, _ string, _ attachment.CreateAttachmentParams) (string, error) {
return "", storeerr.ErrForbidden
}
type fakeAttachmentStoreWithAvatarDownload struct {
fakeAttachmentStore
download attachment.AttachmentDownload
}
func (f *fakeAttachmentStoreWithAvatarDownload) LoadUserAvatar(_ context.Context, _ string) (attachment.AttachmentDownload, error) {
return f.download, nil
}
func (f *fakeAttachmentStoreWithAvatarDownload) LoadAgentAvatar(_ context.Context, _ string, _ string) (attachment.AttachmentDownload, error) {
return f.download, nil
}
func (f *fakeAttachmentStoreWithAvatarDownload) LoadAvatar(_ context.Context, _ string) (attachment.AttachmentDownload, error) {
return f.download, nil
}
// fakeAvatarThumbnailStore drives avatar thumbnail routes: it resolves every
// owner to the same avatar ref and serves one configured blob as the original.
type fakeAvatarThumbnailStore struct {
fakeAttachmentStore
ref attachment.AvatarVariantRef
download attachment.AttachmentDownload
cached attachment.AttachmentVariant
loadCalls int
saveCalls int
saveErr error
savedVariant string
savedBlob []byte
}
func (f *fakeAvatarThumbnailStore) ResolveUserAvatarRef(_ context.Context, _ string) (attachment.AvatarVariantRef, error) {
return f.ref, nil
}
func (f *fakeAvatarThumbnailStore) ResolveAgentAvatarRef(_ context.Context, _ string, _ string) (attachment.AvatarVariantRef, error) {
return f.ref, nil
}
func (f *fakeAvatarThumbnailStore) LoadAvatar(_ context.Context, _ string) (attachment.AttachmentDownload, error) {
f.loadCalls++
return f.download, nil
}
func (f *fakeAvatarThumbnailStore) LoadAvatarThumbnail(_ context.Context, _ string, _ string) (attachment.AttachmentVariant, error) {
if f.cached.Blob != nil {
return f.cached, nil
}
return attachment.AttachmentVariant{}, attachment.ErrNotFound
}
func (f *fakeAvatarThumbnailStore) SaveAvatarThumbnail(_ context.Context, _ string, variant string, gen attachment.GeneratedAvatarThumbnail) (attachment.AttachmentVariant, error) {
f.saveCalls++
f.savedVariant = variant
f.savedBlob = gen.Blob
if f.saveErr != nil {
return attachment.AttachmentVariant{}, f.saveErr
}
return attachment.AttachmentVariant{ContentType: gen.ContentType, Blob: gen.Blob, Width: gen.Width, Height: gen.Height}, nil
}
// testAvatarPNG encodes a width×height PNG avatar with a colored diagonal.
func testAvatarPNG(t *testing.T, width int, height int) []byte {
t.Helper()
img := image.NewRGBA(image.Rect(0, 0, width, height))
for y := 0; y < height; y++ {
for x := 0; x < width; x++ {
img.Set(x, y, color.RGBA{R: uint8(x % 256), G: uint8(y % 256), B: 0x80, A: 0xff})
}
}
var buf bytes.Buffer
require.NoError(t, png.Encode(&buf, img))
return buf.Bytes()
}
type fakeAttachmentStoreWithAttachmentDownload struct {
fakeAttachmentStore
download attachment.AttachmentDownload
memberWorkspaceID string
userWorkspaceID string
}
func (f *fakeAttachmentStoreWithAttachmentDownload) LoadAttachmentForMember(_ context.Context, _ string, workspaceID string, _ string) (attachment.AttachmentDownload, error) {
f.memberWorkspaceID = workspaceID
return f.download, nil
}
func (f *fakeAttachmentStoreWithAttachmentDownload) LoadAttachmentForUser(_ context.Context, _ string, workspaceID string, _ string) (attachment.AttachmentDownload, error) {
f.userWorkspaceID = workspaceID
return f.download, nil
}
type fakeAttachmentHandlerStore interface {
api.SessionAuthenticator
api.WorkspaceResolver
AgentTokenStore
AttachmentStore
}
func newFakeAttachmentHandler(fake fakeAttachmentHandlerStore) *AttachmentHandler {
return &AttachmentHandler{
authStore: fake,
workspaceStore: fake,
agentTokenStore: fake,
store: fake,
}
}
// --- Helper for multipart bodies ---
func multipartBody(t *testing.T, fields map[string]string, fileField, filename string, content []byte) (*bytes.Buffer, string) {
t.Helper()
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
for key, value := range fields {
require.NoError(t, writer.WriteField(key, value))
}
part, err := writer.CreateFormFile(fileField, filename)
require.NoError(t, err)
_, err = part.Write(content)
require.NoError(t, err)
require.NoError(t, writer.Close())
return body, writer.FormDataContentType()
}
// --- AttachmentHandler tests ---
func TestAttachmentHandlerUploadRejectsUnauthenticated(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.POST("/api/v1/workspaces/wsp_1/attachments", handler.Upload)
body, contentType := multipartBody(t, map[string]string{"scopeId": "chn_1"}, "files", "hello.txt", []byte("hello"))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/attachments", body)
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusUnauthorized, recorder.Code)
}
func TestAttachmentHandlerUploadRejectsMissingWorkspace(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.POST("/api/v1/workspaces/wsp_1/attachments", func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}})
handler.Upload(ctx)
})
body, contentType := multipartBody(t, map[string]string{"scopeId": "chn_1"}, "files", "hello.txt", []byte("hello"))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/attachments", body)
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "No workspace selected")
}
func TestAttachmentHandlerUploadsManagedAvatarAsset(t *testing.T) {
t.Parallel()
store := &fakeAttachmentStore{avatarID: "avt_resource"}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.PUT(
"/api/v1/workspaces/:workspaceId/avatar-assets/:ownerKind/:ownerId",
func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}})
api.SetCurrentWorkspace(ctx, workspace.Summary{WorkspaceID: "wsp_1"})
handler.UploadAvatarAsset(ctx)
},
)
tinyPNG, err := base64.StdEncoding.DecodeString(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+kfuoAAAAASUVORK5CYII=",
)
require.NoError(t, err)
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="service.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write(tinyPNG)
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPut,
"/api/v1/workspaces/wsp_1/avatar-assets/service_account/sacct_1",
body,
)
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.JSONEq(t, `{"id":"avt_resource","url":"/api/v1/avatar-assets/avt_resource"}`, recorder.Body.String())
require.Equal(t, "wsp_1", store.avatarWsp)
require.Equal(t, "service_account", store.avatarKind)
require.Equal(t, "sacct_1", store.avatarOwner)
}
func TestAttachmentHandlerUploadsManagedAgentAvatar(t *testing.T) {
t.Parallel()
store := &fakeAttachmentStore{avatarID: "avt_agent"}
agentStore := &fakeAgentAvatarProfileStore{}
handler := newFakeAttachmentHandler(store)
handler.agentStore = agentStore
router := gin.New()
router.PUT(
"/api/v1/workspaces/:workspaceId/agents/:agentId/avatar",
func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}})
api.SetCurrentWorkspace(ctx, workspace.Summary{WorkspaceID: "wsp_1"})
handler.UploadAgentAvatar(ctx)
},
)
tinyPNG, err := base64.StdEncoding.DecodeString(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+kfuoAAAAASUVORK5CYII=",
)
require.NoError(t, err)
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="agent.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write(tinyPNG)
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(
http.MethodPut,
"/api/v1/workspaces/wsp_1/agents/agt_1/avatar",
body,
)
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.JSONEq(t, `{"id":"avt_agent","url":"/api/v1/workspaces/wsp_1/agents/agt_1/avatar?v=avt_agent"}`, recorder.Body.String())
require.Equal(t, 1, agentStore.getCalls)
require.Equal(t, 1, agentStore.updateCalls)
require.Equal(t, "agt_1", agentStore.updatedAgent)
require.Equal(t, "/api/v1/workspaces/wsp_1/agents/agt_1/avatar?v=avt_agent", agentStore.updatedURL)
require.Equal(t, "wsp_1", store.avatarWsp)
require.Equal(t, "agent", store.avatarKind)
require.Equal(t, "agt_1", store.avatarOwner)
}
func TestAttachmentHandlerRejectsAgentAvatarUploadBeforeStoring(t *testing.T) {
t.Parallel()
store := &fakeAttachmentStore{avatarID: "avt_agent"}
agentStore := &fakeAgentAvatarProfileStore{getErr: storeerr.ErrForbidden}
handler := newFakeAttachmentHandler(store)
handler.agentStore = agentStore
router := gin.New()
router.PUT(
"/api/v1/workspaces/:workspaceId/agents/:agentId/avatar",
func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_2"}})
api.SetCurrentWorkspace(ctx, workspace.Summary{WorkspaceID: "wsp_1"})
handler.UploadAgentAvatar(ctx)
},
)
body, contentType := multipartBody(t, nil, "file", "agent.png", []byte("png"))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/workspaces/wsp_1/agents/agt_1/avatar", body)
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusForbidden, recorder.Code)
require.Equal(t, 1, agentStore.getCalls)
require.Zero(t, agentStore.updateCalls)
require.Empty(t, store.avatarOwner)
}
func TestAttachmentHandlerUploadPropagatesStoreError(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentUploadErrorStore{})
router := gin.New()
router.POST("/api/v1/workspaces/wsp_1/attachments", func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}})
api.SetCurrentWorkspace(ctx, workspace.Summary{WorkspaceID: "wsp_1"})
handler.Upload(ctx)
})
body, contentType := multipartBody(t, map[string]string{"scopeId": "chn_1"}, "files", "hello.txt", []byte("hello"))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/attachments", body)
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusForbidden, recorder.Code)
}
func TestAttachmentHandlerUploadRejectsFileTooLarge(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.POST("/api/v1/workspaces/wsp_1/attachments", func(ctx *gin.Context) {
api.SetAuthenticatedUser(ctx, identity.AuthenticatedUser{User: identity.User{ID: "usr_1"}})
api.SetCurrentWorkspace(ctx, workspace.Summary{WorkspaceID: "wsp_1"})
handler.Upload(ctx)
})
oversized := make([]byte, maxMessageAttachmentSizeB+1)
body, contentType := multipartBody(t, map[string]string{"scopeId": "chn_1"}, "files", "big.txt", oversized)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPost, "/api/v1/workspaces/wsp_1/attachments", body)
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "file too large")
}
func TestAttachmentHandlerUploadAvatarRejectsMissingToken(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.PUT("/api/v1/me/avatar", handler.UploadAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusUnauthorized, recorder.Code)
}
func TestAttachmentHandlerUploadAvatarRejectsUnverifiedUserWhenRequired(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
handler.emailVerificationRequired = true
router := gin.New()
router.PUT("/api/v1/me/avatar", handler.UploadAvatar)
body, contentType := multipartBody(t, nil, "file", "avatar.png", []byte("\x89PNG\r\n\x1a\nimage-bytes"))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", body)
request.Header.Set("Authorization", "Bearer human-token")
request.Header.Set(api.HeaderWorkspaceID, "wsp_1")
request.Header.Set("Content-Type", contentType)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusForbidden, recorder.Code)
require.Contains(t, recorder.Body.String(), "email_not_verified")
}
func TestAttachmentHandlerUploadAvatarUsesControlWorkspaceResolver(t *testing.T) {
t.Parallel()
ownedWorkspace := &fakeWorkspaceResolver{err: storeerr.ErrForbidden}
controlWorkspace := &fakeWorkspaceResolver{workspace: workspace.Summary{WorkspaceID: "wsp_control"}}
handler := &AttachmentHandler{
authStore: &fakeAttachmentStore{},
workspaceStore: controlWorkspace,
agentTokenStore: &fakeAttachmentStore{},
store: &fakeAttachmentStore{avatarID: "avt_1"},
emailVerificationRequired: false,
}
router := gin.New()
RegisterControlRoutes(router, Dependencies{
AuthStore: handler.authStore,
WorkspaceStore: ownedWorkspace,
ControlWorkspaceStore: controlWorkspace,
AgentTokenStore: handler.agentTokenStore,
AttachmentStore: handler.store,
})
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="avatar.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write([]byte("\x89PNG\r\n\x1a\nimage-bytes"))
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", body)
request.Header.Set("Authorization", "Bearer human-token")
request.Header.Set(api.HeaderWorkspaceID, "wsp_control")
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, 0, ownedWorkspace.calls)
require.Equal(t, 1, controlWorkspace.calls)
require.Contains(t, recorder.Body.String(), "/api/v1/users/usr_1/avatar")
}
func TestAttachmentHandlerUploadAvatarRejectsForbiddenControlWorkspace(t *testing.T) {
t.Parallel()
controlWorkspace := &fakeWorkspaceResolver{err: storeerr.ErrForbidden}
router := gin.New()
RegisterControlRoutes(router, Dependencies{
AuthStore: &fakeAttachmentStore{},
WorkspaceStore: &fakeAttachmentStore{},
ControlWorkspaceStore: controlWorkspace,
AgentTokenStore: &fakeAttachmentStore{},
AttachmentStore: &fakeAttachmentStore{avatarID: "avt_1"},
})
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="avatar.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write([]byte("\x89PNG\r\n\x1a\nimage-bytes"))
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", body)
request.Header.Set("Authorization", "Bearer human-token")
request.Header.Set(api.HeaderWorkspaceID, "wsp_forbidden")
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusForbidden, recorder.Code)
require.Equal(t, 1, controlWorkspace.calls)
}
func TestAttachmentHandlerUploadAvatarRejectsFileTooLarge(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.PUT("/api/v1/me/avatar", handler.UploadAvatar)
oversized := make([]byte, maxImageUploadSizeB+1)
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="big.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write(oversized)
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", body)
request.Header.Set("Authorization", "Bearer human-token")
request.Header.Set(api.HeaderWorkspaceID, "wsp_1")
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "file too large")
}
func TestAttachmentHandlerUploadAvatarPropagatesStoreError(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAvatarUploadErrorStore{})
router := gin.New()
router.PUT("/api/v1/me/avatar", handler.UploadAvatar)
body := &bytes.Buffer{}
writer := multipart.NewWriter(body)
part, err := writer.CreatePart(textproto.MIMEHeader{
"Content-Disposition": []string{`form-data; name="file"; filename="avatar.png"`},
"Content-Type": []string{"image/png"},
})
require.NoError(t, err)
_, err = part.Write([]byte("\x89PNG\r\n\x1a\nimage-bytes"))
require.NoError(t, err)
require.NoError(t, writer.Close())
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodPut, "/api/v1/me/avatar", body)
request.Header.Set("Authorization", "Bearer human-token")
request.Header.Set(api.HeaderWorkspaceID, "wsp_1")
request.Header.Set("Content-Type", writer.FormDataContentType())
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusForbidden, recorder.Code)
}
func TestAttachmentHandlerGetDeprecatedAvatarObjectRejectsEmptyID(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.GET("/api/avatars/:avatarId", handler.GetDeprecatedAvatarObject)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/avatars/%20", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "avatarId is required")
}
func TestAttachmentHandlerGetDeprecatedAvatarObjectStreamsPayloadWhenPublicURLAlsoPresent(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStoreWithAvatarDownload{
download: attachment.AttachmentDownload{
Filename: "avatar.png",
ContentType: "image/png",
PublicURL: "https://media.agentteam.pingcap.com/agentteam/objects/avatars/wsp_1/avt_1",
Blob: []byte("png-data"),
},
})
router := gin.New()
router.GET("/api/avatars/:avatarId", handler.GetDeprecatedAvatarObject)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/avatars/avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/png", recorder.Header().Get("Content-Type"))
require.Equal(t, "png-data", recorder.Body.String())
require.Empty(t, recorder.Header().Get("Location"))
require.Equal(t, "public, max-age=86400", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetAvatarAssetStreamsPayloadWhenPublicURLAlsoPresent(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStoreWithAvatarDownload{
download: attachment.AttachmentDownload{
Filename: "service-account.svg",
ContentType: "image/svg+xml",
PublicURL: "https://media.agentteam.pingcap.com/agentteam/objects/avatars/wsp_1/avt_1",
Blob: []byte(""),
},
})
router := gin.New()
router.GET("/api/v1/avatar-assets/:avatarId", handler.GetAvatarAsset)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/avatar-assets/avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/svg+xml", recorder.Header().Get("Content-Type"))
require.Equal(t, "", recorder.Body.String())
require.Empty(t, recorder.Header().Get("Location"))
require.Equal(t, "public, max-age=86400", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetUserAvatarRevalidatesUnversionedResource(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.GET("/api/v1/users/:userId/avatar", handler.GetUserAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "no-cache", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetUserAvatarCachesVersionedResource(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.GET("/api/v1/users/:userId/avatar", handler.GetUserAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar?v=avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "public, max-age=86400", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetUserAvatarRedirectsExternalURLWhenNoPayload(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStoreWithAvatarDownload{
download: attachment.AttachmentDownload{
Filename: "avatar",
PublicURL: "https://example.com/avatar.png",
},
})
router := gin.New()
router.GET("/api/v1/users/:userId/avatar", handler.GetUserAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusTemporaryRedirect, recorder.Code)
require.Equal(t, "https://example.com/avatar.png", recorder.Header().Get("Location"))
require.Equal(t, "no-cache", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetAgentAvatarRejectsEmptyAgentID(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStore{})
router := gin.New()
router.GET("/api/v1/workspaces/:workspaceId/agents/:agentId/avatar", handler.GetAgentAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_1/agents/%20/avatar", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "agentId is required")
}
func TestAttachmentHandlerGetAgentAvatarStreamsPayloadWhenPublicURLAlsoPresent(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAttachmentStoreWithAvatarDownload{
download: attachment.AttachmentDownload{
Filename: "agent-avatar.png",
ContentType: "image/png",
PublicURL: "https://media.agentteam.pingcap.com/agentteam/objects/avatars/wsp_1/agt_1",
Blob: []byte("agent-png-data"),
},
})
router := gin.New()
router.GET("/api/v1/workspaces/:workspaceId/agents/:agentId/avatar", handler.GetAgentAvatar)
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/workspaces/wsp_1/agents/agt_1/avatar", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/png", recorder.Header().Get("Content-Type"))
require.Equal(t, "agent-png-data", recorder.Body.String())
require.Empty(t, recorder.Header().Get("Location"))
require.Equal(t, "no-cache", recorder.Header().Get("Cache-Control"))
}
func TestAttachmentHandlerGetUserAvatarThumbnailGeneratesSmallVariant(t *testing.T) {
t.Parallel()
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"},
download: attachment.AttachmentDownload{
Filename: "avatar.png",
ContentType: "image/png",
Blob: testAvatarPNG(t, 200, 300),
},
}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail?v=avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/png", recorder.Header().Get("Content-Type"))
require.Equal(t, "public, max-age=86400", recorder.Header().Get("Cache-Control"))
require.Equal(t, `"avt_1:avatar-cover-128-v1"`, recorder.Header().Get("ETag"))
require.Equal(t, 1, store.loadCalls)
require.Equal(t, 1, store.saveCalls)
require.Equal(t, attachment.AvatarThumbnailSmallCoverVariant, store.savedVariant)
// 200×300 → center-cropped and scaled to the 128px cover tier.
config, _, err := image.DecodeConfig(bytes.NewReader(recorder.Body.Bytes()))
require.NoError(t, err)
require.Equal(t, 128, config.Width)
require.Equal(t, 128, config.Height)
}
func TestAttachmentHandlerGetUserAvatarThumbnailServesGeneratedBytesWhenCacheSaveFails(t *testing.T) {
t.Parallel()
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"},
download: attachment.AttachmentDownload{
Filename: "avatar.png",
ContentType: "image/png",
Blob: testAvatarPNG(t, 200, 300),
},
saveErr: errors.New("object storage write failed"),
}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail?v=avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/png", recorder.Header().Get("Content-Type"))
require.Equal(t, store.savedBlob, recorder.Body.Bytes())
require.NotEmpty(t, recorder.Body.Bytes())
require.Equal(t, 1, store.saveCalls)
}
func TestAttachmentHandlerGetUserAvatarThumbnailGeneratesLargeVariantForSize512(t *testing.T) {
t.Parallel()
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"},
download: attachment.AttachmentDownload{
Filename: "avatar.png",
ContentType: "image/png",
Blob: testAvatarPNG(t, 900, 600),
},
}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail?size=512", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, "image/png", recorder.Header().Get("Content-Type"))
require.Equal(t, "no-cache", recorder.Header().Get("Cache-Control"))
require.Equal(t, attachment.AvatarThumbnailLargeCoverVariant, store.savedVariant)
config, _, err := image.DecodeConfig(bytes.NewReader(recorder.Body.Bytes()))
require.NoError(t, err)
require.Equal(t, 512, config.Width)
require.Equal(t, 512, config.Height)
}
func TestAttachmentHandlerGetUserAvatarThumbnailServesCachedVariantWithoutLoadingOriginal(t *testing.T) {
t.Parallel()
cachedBlob := testAvatarPNG(t, 128, 128)
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"},
cached: attachment.AttachmentVariant{ContentType: "image/png", Blob: cachedBlob},
}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail?v=avt_1", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusOK, recorder.Code)
require.Equal(t, cachedBlob, recorder.Body.Bytes())
require.Zero(t, store.loadCalls)
require.Zero(t, store.saveCalls)
}
func TestAttachmentHandlerGetUserAvatarThumbnailRejectsInvalidSize(t *testing.T) {
t.Parallel()
handler := newFakeAttachmentHandler(&fakeAvatarThumbnailStore{})
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail?size=999", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusBadRequest, recorder.Code)
require.Contains(t, recorder.Body.String(), "size must be 128 or 512")
}
func TestAttachmentHandlerGetUserAvatarThumbnailRedirectsExternalAvatar(t *testing.T) {
t.Parallel()
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{ExternalURL: "https://example.com/avatar.png"},
}
handler := newFakeAttachmentHandler(store)
router := gin.New()
router.GET("/api/v1/users/:userId/avatar/thumbnail", api.HandlePublic(handler.GetUserAvatarThumbnail))
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, "/api/v1/users/usr_1/avatar/thumbnail", nil)
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusTemporaryRedirect, recorder.Code)
require.Equal(t, "/api/v1/users/usr_1/avatar", recorder.Header().Get("Location"))
}
func TestAttachmentHandlerGetUserAvatarThumbnailRedirectsWhenUndecodable(t *testing.T) {
t.Parallel()
store := &fakeAvatarThumbnailStore{
ref: attachment.AvatarVariantRef{AvatarID: "avt_1", WorkspaceID: "wsp_1"},
download: attachment.AttachmentDownload{
Filename: "avatar.svg",
ContentType: "image/svg+xml",
Blob: []byte("