package chat import ( "bytes" "context" "crypto/sha256" "encoding/hex" "encoding/json" "errors" "io" "mime/multipart" "net/http" "net/http/httptest" "net/textproto" "testing" "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/infra/storage" "github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware" ) // ===== upload-specific fakes ===== // uploadRepo wraps fakeRepo and overrides InsertAttachment to capture or fail. type uploadRepo struct { *fakeRepo inserted []InsertAttachmentParams insertErr error storedAtt *Attachment } func (r *uploadRepo) InsertAttachment(_ context.Context, p InsertAttachmentParams) (*Attachment, error) { r.inserted = append(r.inserted, p) if r.insertErr != nil { return nil, r.insertErr } att := &Attachment{ ID: p.ID, UserID: p.UserID, Filename: p.Filename, MimeType: p.MimeType, SizeBytes: p.SizeBytes, SHA256: p.SHA256, StoragePath: p.StoragePath, } r.storedAtt = att return att, nil } // uploadStorage records put/delete and supports a forced putErr. type uploadStorage struct { blobs map[string][]byte putErr error deleteCalls []string } func newUploadStorage() *uploadStorage { return &uploadStorage{blobs: make(map[string][]byte)} } func (s *uploadStorage) Put(_ context.Context, key string, r io.Reader, _ string) (storage.PutResult, error) { if s.putErr != nil { return storage.PutResult{}, s.putErr } data, err := io.ReadAll(r) if err != nil { return storage.PutResult{}, err } s.blobs[key] = data h := sha256.Sum256(data) return storage.PutResult{StoragePath: key, SHA256: hex.EncodeToString(h[:]), Size: int64(len(data))}, nil } func (s *uploadStorage) Get(_ context.Context, key string) (io.ReadCloser, error) { data, ok := s.blobs[key] if !ok { return nil, errors.New("not found: " + key) } return io.NopCloser(bytes.NewReader(data)), nil } func (s *uploadStorage) Delete(_ context.Context, key string) error { s.deleteCalls = append(s.deleteCalls, key) delete(s.blobs, key) return nil } // ===== test helpers ===== // uploadHarness builds a chi router with Auth middleware that resolves any // non-empty bearer token to uid, then mounts the Uploader under POST /uploads. type uploadHarness struct { uid uuid.UUID repo *uploadRepo stor *uploadStorage up *Uploader r chi.Router } func newUploadHarness(t *testing.T, mimes []string, maxFile int64) *uploadHarness { t.Helper() uid := uuid.Must(uuid.NewV7()) repo := &uploadRepo{fakeRepo: newFakeRepo()} stor := newUploadStorage() up := NewUploader(repo, stor, mimes, maxFile) resolver := &fakeSessionResolver{uid: uid} r := chi.NewRouter() r.Use(middleware.RequestID) r.Use(middleware.Auth(resolver, middleware.AuthOptions{})) r.Post("/uploads", up.Handle) return &uploadHarness{uid: uid, repo: repo, stor: stor, up: up, r: r} } // buildMultipart constructs a multipart/form-data body with one "file" part. // declaredMime is what the part's Content-Type header advertises. func buildMultipart(t *testing.T, filename, declaredMime string, data []byte) (body *bytes.Buffer, contentType string) { t.Helper() body = &bytes.Buffer{} mw := multipart.NewWriter(body) hdr := textproto.MIMEHeader{} hdr.Set("Content-Disposition", `form-data; name="file"; filename="`+filename+`"`) hdr.Set("Content-Type", declaredMime) part, err := mw.CreatePart(hdr) require.NoError(t, err) _, err = part.Write(data) require.NoError(t, err) require.NoError(t, mw.Close()) return body, mw.FormDataContentType() } func newUploadRequest(t *testing.T, body *bytes.Buffer, contentType string, withAuth bool) *http.Request { t.Helper() req := httptest.NewRequest(http.MethodPost, "/uploads", body) req.Header.Set("Content-Type", contentType) if withAuth { req.Header.Set("Authorization", "Bearer test-token") } return req } // ===== tests ===== func TestUploader_HappyPath_ReturnsAttachmentMetadata(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 1<<20) data := []byte("hello world") body, ct := buildMultipart(t, "greeting.txt", "text/plain", data) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String()) var resp AttachmentDTO require.NoError(t, json.Unmarshal(w.Body.Bytes(), &resp)) assert.Equal(t, "greeting.txt", resp.Filename) assert.Equal(t, "text/plain", resp.MimeType) assert.Equal(t, int64(len(data)), resp.SizeBytes) expectedSHA := sha256.Sum256(data) assert.Equal(t, hex.EncodeToString(expectedSHA[:]), resp.SHA256) require.Len(t, h.repo.inserted, 1) got := h.repo.inserted[0] assert.Equal(t, h.uid, got.UserID) assert.Equal(t, "greeting.txt", got.Filename) } func TestUploader_FileTooLarge_413(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 5) body, ct := buildMultipart(t, "big.txt", "text/plain", []byte("more than five bytes")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusRequestEntityTooLarge, w.Code, "body: %s", w.Body.String()) assert.Empty(t, h.repo.inserted, "no DB row should be inserted") } func TestUploader_UnsupportedMime_415(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 1<<20) body, ct := buildMultipart(t, "evil.exe", "application/octet-stream", []byte("MZ...")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusUnsupportedMediaType, w.Code) assert.Empty(t, h.repo.inserted) assert.Empty(t, h.stor.blobs, "blob must not be stored when mime is rejected") } func TestUploader_StoragePutFails_NoOrphanRow(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 1<<20) h.stor.putErr = errors.New("disk full") body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusBadGateway, w.Code) assert.Empty(t, h.repo.inserted, "no DB insert when storage fails") assert.Empty(t, h.stor.deleteCalls, "no delete when nothing was stored") } func TestUploader_DBInsertFails_DeletesBlob(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 1<<20) h.repo.insertErr = errs.New(errs.CodeInternal, "db down") body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusInternalServerError, w.Code) require.Len(t, h.repo.inserted, 1) require.Len(t, h.stor.deleteCalls, 1, "blob must be cleaned up after DB failure") assert.Equal(t, h.repo.inserted[0].StoragePath, h.stor.deleteCalls[0]) } func TestUploader_NoAuth_401(t *testing.T) { h := newUploadHarness(t, []string{"text/plain"}, 1<<20) body, ct := buildMultipart(t, "x.txt", "text/plain", []byte("hi")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, false)) require.Equal(t, http.StatusUnauthorized, w.Code) assert.Empty(t, h.repo.inserted) } func TestUploader_TextWildcard_Allows(t *testing.T) { h := newUploadHarness(t, []string{"text/*"}, 1<<20) body, ct := buildMultipart(t, "page.html", "text/html", []byte("
ok
")) w := httptest.NewRecorder() h.r.ServeHTTP(w, newUploadRequest(t, body, ct, true)) require.Equal(t, http.StatusOK, w.Code, "body: %s", w.Body.String()) require.Len(t, h.repo.inserted, 1) assert.Equal(t, "text/html", h.repo.inserted[0].MimeType) }