mirror of
https://github.com/langgenius/dify.git
synced 2026-09-05 00:31:19 +08:00
251 lines
8.4 KiB
Go
251 lines
8.4 KiB
Go
package agentcli
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
type fakeFileUploadClient struct {
|
|
forFrontend bool
|
|
downloadRequestCall int
|
|
uploadResponse []byte
|
|
calls []string
|
|
filename string
|
|
mimetype string
|
|
uploadURL string
|
|
uploadedBytes []byte
|
|
downloadReference string
|
|
}
|
|
|
|
func (f *fakeFileUploadClient) CreateFileUploadURL(_ context.Context, filename, mimetype string) (string, error) {
|
|
f.calls = append(f.calls, "upload-request")
|
|
f.filename = filename
|
|
f.mimetype = mimetype
|
|
return "https://sandbox-files.example.com/files/upload/for-plugin?sign=1", nil
|
|
}
|
|
|
|
func (f *fakeFileUploadClient) UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error) {
|
|
f.calls = append(f.calls, "multipart-upload")
|
|
f.uploadURL = uploadURL
|
|
f.filename = filename
|
|
f.mimetype = mimetype
|
|
f.uploadedBytes, _ = os.ReadFile(filePath)
|
|
if f.uploadResponse != nil {
|
|
return f.uploadResponse, nil
|
|
}
|
|
return []byte(`{"reference":"dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ=="}`), nil
|
|
}
|
|
|
|
func (f *fakeFileUploadClient) CreateFileDownloadURL(
|
|
_ context.Context,
|
|
_ string,
|
|
reference, _ *string,
|
|
forFrontend bool,
|
|
) (*FileDownloadResponse, error) {
|
|
f.calls = append(f.calls, "download-request")
|
|
f.forFrontend = forFrontend
|
|
f.downloadRequestCall++
|
|
if reference != nil {
|
|
f.downloadReference = *reference
|
|
}
|
|
return &FileDownloadResponse{
|
|
Filename: "report.pdf",
|
|
MimeType: "application/pdf",
|
|
Size: 123,
|
|
DownloadURL: "/files/tools/report.pdf?sign=2",
|
|
}, nil
|
|
}
|
|
|
|
func TestRunFileUploadReturnsFrontendDisplayURL(t *testing.T) {
|
|
filePath := t.TempDir() + "/report.pdf"
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
client := &fakeFileUploadClient{}
|
|
var output bytes.Buffer
|
|
if err := runFileUpload(client, filePath, false, &output); err != nil {
|
|
t.Fatalf("run file upload: %v", err)
|
|
}
|
|
|
|
if !client.forFrontend {
|
|
t.Fatal("download request did not select frontend display URL")
|
|
}
|
|
if got, want := strings.Join(client.calls, ","), "upload-request,multipart-upload,download-request"; got != want {
|
|
t.Fatalf("call order = %s, want %s", got, want)
|
|
}
|
|
if client.filename != "report.pdf" || client.mimetype != "application/pdf" {
|
|
t.Fatalf("upload metadata = (%q, %q), want report.pdf/application/pdf", client.filename, client.mimetype)
|
|
}
|
|
if client.uploadURL != "https://sandbox-files.example.com/files/upload/for-plugin?sign=1" {
|
|
t.Fatalf("upload URL = %q", client.uploadURL)
|
|
}
|
|
if string(client.uploadedBytes) != "report" {
|
|
t.Fatalf("uploaded bytes = %q, want report", client.uploadedBytes)
|
|
}
|
|
if client.downloadReference != "dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ==" {
|
|
t.Fatalf("download reference = %q", client.downloadReference)
|
|
}
|
|
got := strings.TrimSpace(output.String())
|
|
want := `{"transfer_method":"tool_file","reference":"dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ==","public_download_url":"/files/tools/report.pdf?sign=2"}`
|
|
if got != want {
|
|
t.Fatalf("output = %s, want %s", got, want)
|
|
}
|
|
}
|
|
|
|
func TestRunFileUploadWithoutDownloadLinkReturnsOnlyCanonicalMapping(t *testing.T) {
|
|
filePath := filepath.Join(t.TempDir(), "report.pdf")
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
client := &fakeFileUploadClient{}
|
|
var output bytes.Buffer
|
|
if err := runFileUpload(client, filePath, true, &output); err != nil {
|
|
t.Fatalf("run file upload: %v", err)
|
|
}
|
|
|
|
if client.downloadRequestCall != 0 {
|
|
t.Fatalf("download request calls = %d, want 0", client.downloadRequestCall)
|
|
}
|
|
if got, want := strings.Join(client.calls, ","), "upload-request,multipart-upload"; got != want {
|
|
t.Fatalf("call order = %s, want %s", got, want)
|
|
}
|
|
got := strings.TrimSpace(output.String())
|
|
want := `{"transfer_method":"tool_file","reference":"dify-file-ref:eyJyZWNvcmRfaWQiOiJ0b29sLTEifQ=="}`
|
|
if got != want {
|
|
t.Fatalf("output = %s, want %s", got, want)
|
|
}
|
|
}
|
|
|
|
func TestRunFileUploadDefaultAcceptsLegacyNonemptyReference(t *testing.T) {
|
|
filePath := filepath.Join(t.TempDir(), "report.pdf")
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":"raw-id"}`)}
|
|
var output bytes.Buffer
|
|
err := runFileUpload(client, filePath, false, &output)
|
|
if err != nil {
|
|
t.Fatalf("run file upload: %v", err)
|
|
}
|
|
if client.downloadRequestCall != 1 || client.downloadReference != "raw-id" {
|
|
t.Fatalf("download request = (%d, %q), want legacy reference", client.downloadRequestCall, client.downloadReference)
|
|
}
|
|
if !strings.Contains(output.String(), `"reference":"raw-id"`) {
|
|
t.Fatalf("output = %q, want legacy reference", output.String())
|
|
}
|
|
}
|
|
|
|
func TestRunFileUploadWithoutDownloadLinkRejectsNonCanonicalReference(t *testing.T) {
|
|
filePath := filepath.Join(t.TempDir(), "report.pdf")
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":"raw-id"}`)}
|
|
var output bytes.Buffer
|
|
err := runFileUpload(client, filePath, true, &output)
|
|
if err == nil || !strings.Contains(err.Error(), "invalid reference") {
|
|
t.Fatalf("error = %v, want invalid reference", err)
|
|
}
|
|
if client.downloadRequestCall != 0 {
|
|
t.Fatalf("download request calls = %d, want 0", client.downloadRequestCall)
|
|
}
|
|
if output.Len() != 0 {
|
|
t.Fatalf("output = %q, want empty", output.String())
|
|
}
|
|
}
|
|
|
|
func TestRunFileUploadRejectsMissingReferenceInBothModes(t *testing.T) {
|
|
filePath := filepath.Join(t.TempDir(), "report.pdf")
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
for _, noDownloadLink := range []bool{false, true} {
|
|
client := &fakeFileUploadClient{uploadResponse: []byte(`{"reference":""}`)}
|
|
var output bytes.Buffer
|
|
err := runFileUpload(client, filePath, noDownloadLink, &output)
|
|
if err == nil || !strings.Contains(err.Error(), "missing reference") {
|
|
t.Fatalf("noDownloadLink=%t error = %v, want missing reference", noDownloadLink, err)
|
|
}
|
|
if client.downloadRequestCall != 0 {
|
|
t.Fatalf("noDownloadLink=%t download request calls = %d, want 0", noDownloadLink, client.downloadRequestCall)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestRunFileUploadRejectsInvalidUploadResponseBeforeDownloadRequest(t *testing.T) {
|
|
filePath := filepath.Join(t.TempDir(), "report.pdf")
|
|
if err := os.WriteFile(filePath, []byte("report"), 0o600); err != nil {
|
|
t.Fatalf("write fixture: %v", err)
|
|
}
|
|
|
|
client := &fakeFileUploadClient{uploadResponse: []byte("not-json")}
|
|
var output bytes.Buffer
|
|
err := runFileUpload(client, filePath, true, &output)
|
|
if err == nil || !strings.Contains(err.Error(), "parse upload result") {
|
|
t.Fatalf("error = %v, want parse upload result failure", err)
|
|
}
|
|
if client.downloadRequestCall != 0 {
|
|
t.Fatalf("download request calls = %d, want 0", client.downloadRequestCall)
|
|
}
|
|
}
|
|
|
|
func TestRunFileDownloadRequestsSandboxURLAndWritesFile(t *testing.T) {
|
|
var requestPayload map[string]json.RawMessage
|
|
var server *httptest.Server
|
|
server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
switch r.URL.Path {
|
|
case "/agent-stub/files/download-request":
|
|
if err := json.NewDecoder(r.Body).Decode(&requestPayload); err != nil {
|
|
t.Errorf("decode download request: %v", err)
|
|
http.Error(w, "bad request", http.StatusBadRequest)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{"filename":"report.pdf","mime_type":"application/pdf","size":6,"download_url":"` + server.URL + `/files/report.pdf"}`))
|
|
case "/files/report.pdf":
|
|
_, _ = w.Write([]byte("report"))
|
|
default:
|
|
http.NotFound(w, r)
|
|
}
|
|
}))
|
|
defer server.Close()
|
|
|
|
targetDir := t.TempDir()
|
|
err := RunFileDownload(
|
|
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "test-token"},
|
|
"tool_file",
|
|
"dify-file-ref:canonical",
|
|
targetDir,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("run file download: %v", err)
|
|
}
|
|
|
|
var forFrontend bool
|
|
if err := json.Unmarshal(requestPayload["for_frontend"], &forFrontend); err != nil {
|
|
t.Fatalf("decode for_frontend: %v", err)
|
|
}
|
|
if forFrontend {
|
|
t.Fatal("download request selected a frontend URL")
|
|
}
|
|
data, err := os.ReadFile(filepath.Join(targetDir, "report.pdf"))
|
|
if err != nil {
|
|
t.Fatalf("read downloaded file: %v", err)
|
|
}
|
|
if string(data) != "report" {
|
|
t.Fatalf("downloaded file = %q, want report", data)
|
|
}
|
|
}
|