You've already forked agentic-coding-workflow
232 lines
7.5 KiB
Go
232 lines
7.5 KiB
Go
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("<p>ok</p>"))
|
|
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)
|
|
}
|