starnet/file_upload_test.go
2026-03-08 20:19:40 +08:00

173 lines
4.2 KiB
Go

package starnet
import (
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
)
func TestRequestAddFileStream(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
err := r.ParseMultipartForm(10 << 20) // 10 MB
if err != nil {
t.Fatalf("ParseMultipartForm error: %v", err)
}
file, header, err := r.FormFile("file")
if err != nil {
t.Fatalf("FormFile error: %v", err)
}
defer file.Close()
content, _ := io.ReadAll(file)
w.Write([]byte(header.Filename + ":" + string(content)))
}))
defer server.Close()
fileContent := "test file content"
reader := strings.NewReader(fileContent)
req := NewSimpleRequest(server.URL, "POST").
AddFileStream("file", "test.txt", int64(len(fileContent)), reader)
resp, err := req.Do()
if err != nil {
t.Fatalf("Do() error: %v", err)
}
defer resp.Close()
body, _ := resp.Body().String()
expected := "test.txt:" + fileContent
if body != expected {
t.Errorf("Body = %v; want %v", body, expected)
}
}
func TestRequestAddFileWithFormData(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
err := r.ParseMultipartForm(10 << 20)
if err != nil {
t.Fatalf("ParseMultipartForm error: %v", err)
}
// Check form field
name := r.FormValue("name")
if name != "John" {
t.Errorf("name = %v; want John", name)
}
// Check file
file, header, err := r.FormFile("file")
if err != nil {
t.Fatalf("FormFile error: %v", err)
}
defer file.Close()
w.Write([]byte("OK:" + header.Filename))
}))
defer server.Close()
fileContent := "file data"
reader := strings.NewReader(fileContent)
req := NewSimpleRequest(server.URL, "POST").
AddFormData("name", "John").
AddFileStream("file", "document.txt", int64(len(fileContent)), reader)
resp, err := req.Do()
if err != nil {
t.Fatalf("Do() error: %v", err)
}
defer resp.Close()
body, _ := resp.Body().String()
if !strings.Contains(body, "document.txt") {
t.Errorf("Body should contain filename, got: %v", body)
}
}
func TestRequestUploadProgress(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
r.ParseMultipartForm(10 << 20)
w.WriteHeader(http.StatusOK)
}))
defer server.Close()
progressCalled := false
var lastUploaded int64
fileContent := strings.Repeat("a", 1024*10) // 10KB
reader := strings.NewReader(fileContent)
req := NewSimpleRequest(server.URL, "POST").
SetUploadProgress(func(filename string, uploaded, total int64) {
progressCalled = true
lastUploaded = uploaded
if filename != "test.txt" {
t.Errorf("filename = %v; want test.txt", filename)
}
}).
AddFileStream("file", "test.txt", int64(len(fileContent)), reader)
resp, err := req.Do()
if err != nil {
t.Fatalf("Do() error: %v", err)
}
defer resp.Close()
if !progressCalled {
t.Error("Progress callback was not called")
}
if lastUploaded != int64(len(fileContent)) {
t.Errorf("lastUploaded = %v; want %v", lastUploaded, len(fileContent))
}
}
// TestRequestAddFileFromDisk tests uploading a real file from disk
func TestRequestAddFileFromDisk(t *testing.T) {
// Create a temporary file
tmpDir := t.TempDir()
tmpFile := filepath.Join(tmpDir, "test.txt")
fileContent := []byte("test file content from disk")
err := os.WriteFile(tmpFile, fileContent, 0644)
if err != nil {
t.Fatalf("WriteFile error: %v", err)
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
err := r.ParseMultipartForm(10 << 20)
if err != nil {
t.Fatalf("ParseMultipartForm error: %v", err)
}
file, header, err := r.FormFile("file")
if err != nil {
t.Fatalf("FormFile error: %v", err)
}
defer file.Close()
content, _ := io.ReadAll(file)
w.Write([]byte(header.Filename + ":" + string(content)))
}))
defer server.Close()
req := NewSimpleRequest(server.URL, "POST").AddFile("file", tmpFile)
resp, err := req.Do()
if err != nil {
t.Fatalf("Do() error: %v", err)
}
defer resp.Close()
body, _ := resp.Body().String()
if !strings.Contains(body, string(fileContent)) {
t.Errorf("Body should contain file content, got: %v", body)
}
}