mirror of
https://github.com/langgenius/dify.git
synced 2026-07-20 09:38:32 +08:00
115 lines
3.1 KiB
Go
115 lines
3.1 KiB
Go
package stubclient
|
|
|
|
import (
|
|
"context"
|
|
"net"
|
|
"testing"
|
|
|
|
"google.golang.org/grpc"
|
|
|
|
stubv1 "github.com/langgenius/dify/dify-agent-runtime/gen/dify/agent/stub/v1"
|
|
)
|
|
|
|
// fakeServer implements the AgentStubService for testing.
|
|
type fakeServer struct {
|
|
stubv1.UnimplementedAgentStubServiceServer
|
|
}
|
|
|
|
func (f *fakeServer) Connect(_ context.Context, req *stubv1.ConnectRequest) (*stubv1.ConnectResponse, error) {
|
|
return &stubv1.ConnectResponse{
|
|
ConnectionId: "conn-123",
|
|
Status: "connected",
|
|
}, nil
|
|
}
|
|
|
|
func (f *fakeServer) CreateFileUploadRequest(_ context.Context, req *stubv1.FileUploadRequest) (*stubv1.FileUploadResponse, error) {
|
|
return &stubv1.FileUploadResponse{
|
|
UploadUrl: "https://upload.example.com/" + req.GetFilename(),
|
|
}, nil
|
|
}
|
|
|
|
func (f *fakeServer) CreateFileDownloadRequest(_ context.Context, req *stubv1.FileDownloadRequest) (*stubv1.FileDownloadResponse, error) {
|
|
return &stubv1.FileDownloadResponse{
|
|
Filename: "downloaded.txt",
|
|
MimeType: strPtr("text/plain"),
|
|
Size: 1024,
|
|
DownloadUrl: "https://download.example.com/file",
|
|
}, nil
|
|
}
|
|
|
|
func strPtr(s string) *string { return &s }
|
|
|
|
func startFakeServer(t *testing.T) string {
|
|
t.Helper()
|
|
lis, err := net.Listen("tcp", "127.0.0.1:0")
|
|
if err != nil {
|
|
t.Fatalf("failed to listen: %v", err)
|
|
}
|
|
srv := grpc.NewServer()
|
|
stubv1.RegisterAgentStubServiceServer(srv, &fakeServer{})
|
|
go func() { _ = srv.Serve(lis) }()
|
|
t.Cleanup(srv.Stop)
|
|
return lis.Addr().String()
|
|
}
|
|
|
|
func TestConnect(t *testing.T) {
|
|
addr := startFakeServer(t)
|
|
c, err := Dial(addr)
|
|
if err != nil {
|
|
t.Fatalf("Dial failed: %v", err)
|
|
}
|
|
defer func() { _ = c.Close() }()
|
|
|
|
result, err := c.Connect(context.Background(), 1, []string{"hello"}, `{"key":"val"}`)
|
|
if err != nil {
|
|
t.Fatalf("Connect failed: %v", err)
|
|
}
|
|
if result.ConnectionID != "conn-123" {
|
|
t.Errorf("expected conn-123, got %s", result.ConnectionID)
|
|
}
|
|
if result.Status != "connected" {
|
|
t.Errorf("expected connected, got %s", result.Status)
|
|
}
|
|
}
|
|
|
|
func TestCreateFileUpload(t *testing.T) {
|
|
addr := startFakeServer(t)
|
|
c, err := Dial(addr)
|
|
if err != nil {
|
|
t.Fatalf("Dial failed: %v", err)
|
|
}
|
|
defer func() { _ = c.Close() }()
|
|
|
|
result, err := c.CreateFileUpload(context.Background(), "test.txt", "text/plain")
|
|
if err != nil {
|
|
t.Fatalf("CreateFileUpload failed: %v", err)
|
|
}
|
|
if result.UploadURL != "https://upload.example.com/test.txt" {
|
|
t.Errorf("unexpected upload URL: %s", result.UploadURL)
|
|
}
|
|
}
|
|
|
|
func TestCreateFileDownload(t *testing.T) {
|
|
addr := startFakeServer(t)
|
|
c, err := Dial(addr)
|
|
if err != nil {
|
|
t.Fatalf("Dial failed: %v", err)
|
|
}
|
|
defer func() { _ = c.Close() }()
|
|
|
|
ref := "dify-file-ref:abc123"
|
|
result, err := c.CreateFileDownload(context.Background(), "local_file", &ref, nil, false)
|
|
if err != nil {
|
|
t.Fatalf("CreateFileDownload failed: %v", err)
|
|
}
|
|
if result.Filename != "downloaded.txt" {
|
|
t.Errorf("unexpected filename: %s", result.Filename)
|
|
}
|
|
if result.Size != 1024 {
|
|
t.Errorf("unexpected size: %d", result.Size)
|
|
}
|
|
if result.DownloadURL != "https://download.example.com/file" {
|
|
t.Errorf("unexpected download URL: %s", result.DownloadURL)
|
|
}
|
|
}
|