116 lines
3.6 KiB
Go
116 lines
3.6 KiB
Go
package handler
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/wmms/wmms/internal/service"
|
|
"github.com/wmms/wmms/pkg/response"
|
|
)
|
|
|
|
type countedImportBody struct {
|
|
io.Reader
|
|
read int64
|
|
}
|
|
|
|
func (body *countedImportBody) Read(p []byte) (int, error) {
|
|
n, err := body.Reader.Read(p)
|
|
body.read += int64(n)
|
|
return n, err
|
|
}
|
|
|
|
func (*countedImportBody) Close() error { return nil }
|
|
|
|
func TestMaterialImportRejectsLargeBodyBeforeReadingItAll(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
var payload bytes.Buffer
|
|
writer := multipart.NewWriter(&payload)
|
|
part, err := writer.CreateFormFile("file", "物料.xlsx")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := part.Write(bytes.Repeat([]byte("x"), 8<<20)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, contentLength := range []int64{int64(payload.Len()), -1} {
|
|
t.Run(fmt.Sprintf("content-length=%d", contentLength), func(t *testing.T) {
|
|
body := &countedImportBody{Reader: bytes.NewReader(payload.Bytes())}
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Set("organization_id", int64(1))
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/material-import", nil)
|
|
c.Request.Body = body
|
|
c.Request.ContentLength = contentLength
|
|
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
(&Handler{}).ImportMaterials(c)
|
|
if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), "1MB") {
|
|
t.Fatalf("status=%d, body=%s", recorder.Code, recorder.Body.String())
|
|
}
|
|
if body.read > service.MaterialImportMaxBytes+(64<<10)+1 {
|
|
t.Fatalf("超大请求已读取 %d 字节，应在解析前限制读取", body.read)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestMaterialImportFileSizeBoundaryAllowsMultipartOverhead(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
for _, size := range []int{service.MaterialImportMaxBytes, service.MaterialImportMaxBytes + 1} {
|
|
t.Run(fmt.Sprintf("file-size=%d", size), func(t *testing.T) {
|
|
var payload bytes.Buffer
|
|
writer := multipart.NewWriter(&payload)
|
|
part, err := writer.CreateFormFile("file", "物料.xlsx")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err := part.Write(bytes.Repeat([]byte("x"), size)); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
c.Set("organization_id", int64(1))
|
|
c.Request = httptest.NewRequest(http.MethodPost, "/material-import", bytes.NewReader(payload.Bytes()))
|
|
c.Request.Header.Set("Content-Type", writer.FormDataContentType())
|
|
(&Handler{materialImportSvc: service.NewMaterialImportService()}).ImportMaterials(c)
|
|
want := "模板格式不正确"
|
|
if size > service.MaterialImportMaxBytes {
|
|
want = "文件超过 1MB 上限"
|
|
}
|
|
if recorder.Code != http.StatusBadRequest || !strings.Contains(recorder.Body.String(), want) {
|
|
t.Fatalf("status=%d, body=%s, want %s", recorder.Code, recorder.Body.String(), want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestErrorWithDataKeepsPayload(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
recorder := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(recorder)
|
|
|
|
response.ErrorWithData(c, http.StatusBadRequest, "有 1 行校验未通过", map[string]interface{}{
|
|
"total_rows": 1, "failed_rows": 1,
|
|
})
|
|
|
|
if recorder.Code != http.StatusBadRequest {
|
|
t.Fatalf("status = %d, want 400", recorder.Code)
|
|
}
|
|
body := recorder.Body.String()
|
|
if !strings.Contains(body, "有 1 行校验未通过") || !strings.Contains(body, `"failed_rows":1`) {
|
|
t.Fatalf("body = %s", body)
|
|
}
|
|
}
|