Files
NianAIGC/backend/internal/assets/remote_test.go

43 lines
1.3 KiB
Go

package assets
import (
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"testing"
)
func TestHTTPRemoteFetcherRestrictsProtocolAndSize(t *testing.T) {
fetcher := NewHTTPRemoteFetcher(http.DefaultClient, 4)
if _, err := fetcher.Fetch(context.Background(), "file:///etc/passwd"); !errors.Is(err, ErrRemoteProtocol) {
t.Fatalf("file protocol error = %v", err)
}
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
_, _ = io.WriteString(w, "12345")
}))
defer server.Close()
if _, err := fetcher.Fetch(context.Background(), server.URL); !errors.Is(err, ErrRemoteTooLarge) {
t.Fatalf("oversize error = %v", err)
}
}
func TestHTTPRemoteFetcherReturnsBoundedBody(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "image/png")
_, _ = io.WriteString(w, "1234")
}))
defer server.Close()
blob, err := NewHTTPRemoteFetcher(http.DefaultClient, 4).Fetch(context.Background(), server.URL)
if err != nil {
t.Fatal(err)
}
defer blob.Body.Close()
body, _ := io.ReadAll(blob.Body)
if string(body) != "1234" || blob.Size != 4 || blob.ContentType != "image/png" {
t.Fatalf("blob = %#v body=%q", blob, body)
}
}