package service import ( "bytes" "context" "errors" "io" "net/http" "net/http/httptest" "strings" "testing" "time" "github.com/glebarez/sqlite" "gorm.io/gorm" "ragflow/internal/dao" "ragflow/internal/entity" "ragflow/internal/storage" ) // fakeStorage mocks storage.Storage for testing DownloadAgentFile. type fakeStorage struct { lastBucket string lastFnm string blob []byte err error exists bool getCalls int } func testFileService() *FileService { return &FileService{ fileDAO: dao.NewFileDAO(), file2DocumentDAO: dao.NewFile2DocumentDAO(), documentService: &DocumentService{}, } } func (f *fakeStorage) Health() bool { return true } func (f *fakeStorage) Put(bucket, fnm string, binary []byte, tenantID ...string) error { f.lastBucket = bucket f.lastFnm = fnm f.blob = binary f.exists = true return f.err } func (f *fakeStorage) Get(bucket, fnm string, tenantID ...string) ([]byte, error) { f.getCalls++ f.lastBucket = bucket f.lastFnm = fnm return f.blob, f.err } func (f *fakeStorage) Remove(bucket, fnm string, tenantID ...string) error { panic("not implemented in fakeStorage") } func (f *fakeStorage) ObjExist(bucket, fnm string, tenantID ...string) bool { return f.exists && f.lastBucket == bucket && f.lastFnm == fnm } func (f *fakeStorage) GetPresignedURL(bucket, fnm string, expires time.Duration, tenantID ...string) (string, error) { panic("not implemented in fakeStorage") } func (f *fakeStorage) BucketExists(bucket string) bool { panic("not implemented in fakeStorage") } func (f *fakeStorage) ListObjects(bucket string, tenantID ...string) ([]string, error) { panic("not implemented in fakeStorage") } func (f *fakeStorage) RemoveBucket(bucket string) error { panic("not implemented in fakeStorage") } func (f *fakeStorage) Copy(srcBucket, srcPath, destBucket, destPath string) bool { panic("not implemented in fakeStorage") } func (f *fakeStorage) Move(srcBucket, srcPath, destBucket, destPath string) bool { panic("not implemented in fakeStorage") } func (f *fakeStorage) Close() error { return nil } func setupFileContentPermissionDB(t *testing.T, accessible bool) { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ TranslateError: true, }) if err != nil { t.Fatalf("failed to open sqlite: %v", err) } if err := db.AutoMigrate( &entity.File{}, &entity.File2Document{}, &entity.Document{}, &entity.Knowledgebase{}, &entity.Tenant{}, &entity.UserTenant{}, ); err != nil { t.Fatalf("failed to migrate file content tables: %v", err) } location := "doc.txt" if err := db.Create(&entity.File{ ID: "file-1", ParentID: "bucket-1", TenantID: "owner-user", Name: "doc.txt", Location: &location, Type: "file", }).Error; err != nil { t.Fatalf("insert file: %v", err) } if err := db.Create(&entity.Document{ ID: "doc-1", KbID: "kb-owner", ParserID: "naive", ParserConfig: entity.JSONMap{}, Status: sptr(string(entity.StatusValid)), }).Error; err != nil { t.Fatalf("insert document: %v", err) } fileID := "file-1" docID := "doc-1" if err := db.Create(&entity.File2Document{ ID: "f2d-1", FileID: &fileID, DocumentID: &docID, }).Error; err != nil { t.Fatalf("insert file2document: %v", err) } if err := db.Create(&entity.Knowledgebase{ ID: "kb-owner", TenantID: "tenant-owner", Name: "owner-kb", EmbdID: "embd-1", CreatedBy: "owner-user", Permission: string(entity.TenantPermissionTeam), Status: sptr(string(entity.StatusValid)), }).Error; err != nil { t.Fatalf("insert knowledgebase: %v", err) } if err := db.Create(&entity.Tenant{ ID: "tenant-owner", LLMID: "llm-1", EmbdID: "embd-1", ASRID: "asr-1", Status: sptr(string(entity.StatusValid)), }).Error; err != nil { t.Fatalf("insert tenant: %v", err) } if accessible { if err := db.Create(&entity.UserTenant{ ID: "ut-user-1", UserID: "user-1", TenantID: "tenant-owner", Role: "normal", Status: sptr(string(entity.StatusValid)), }).Error; err != nil { t.Fatalf("insert user_tenant: %v", err) } } orig := dao.DB dao.DB = db t.Cleanup(func() { dao.DB = orig }) } func TestFileService_GetFileContents_NotAccessible(t *testing.T) { setupFileContentPermissionDB(t, false) mockStorage := &fakeStorage{blob: []byte("secret")} factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() texts, images, err := svc.GetFileContents("user-1", []map[string]interface{}{{"id": "file-1"}}, false) if err == nil { t.Fatal("expected authorization error") } if err.Error() != "No authorization." { t.Fatalf("unexpected error: %v", err) } if len(texts) != 0 || len(images) != 0 { t.Fatalf("expected no content, got texts=%v images=%v", texts, images) } if mockStorage.getCalls != 0 { t.Fatalf("storage should not be read without permission, got %d calls", mockStorage.getCalls) } } func TestFileService_GetFileContents_Accessible(t *testing.T) { setupFileContentPermissionDB(t, true) mockStorage := &fakeStorage{blob: []byte("allowed content")} factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() texts, images, err := svc.GetFileContents("user-1", []map[string]interface{}{{"id": "file-1"}}, false) if err != nil { t.Fatalf("GetFileContents failed: %v", err) } if len(images) != 0 { t.Fatalf("expected no images, got %v", images) } if len(texts) != 1 || !strings.Contains(texts[0], "allowed content") { t.Fatalf("unexpected texts: %v", texts) } if mockStorage.getCalls != 1 { t.Fatalf("storage get calls = %d, want 1", mockStorage.getCalls) } if mockStorage.lastBucket != "bucket-1" || mockStorage.lastFnm != "doc.txt" { t.Fatalf("storage read %s/%s, want bucket-1/doc.txt", mockStorage.lastBucket, mockStorage.lastFnm) } } func TestFileService_DownloadAgentFile_Success(t *testing.T) { // Setup mock storage expectedBlob := []byte("fake file content") mockStorage := &fakeStorage{ blob: expectedBlob, err: nil, } factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() tenantID := "tenant123" location := "file-abc.txt" blob, err := svc.DownloadAgentFile(tenantID, location) if err != nil { t.Fatalf("expected no error, got %v", err) } if mockStorage.lastBucket != "tenant123-downloads" { t.Errorf("expected bucket 'tenant123-downloads', got %q", mockStorage.lastBucket) } if mockStorage.lastFnm != location { t.Errorf("expected fnm %q, got %q", location, mockStorage.lastFnm) } if !bytes.Equal(blob, expectedBlob) { t.Errorf("expected blob %v, got %v", expectedBlob, blob) } } func TestFileService_DownloadAgentFile_Error(t *testing.T) { // Setup mock storage expectedErr := errors.New("not found") mockStorage := &fakeStorage{ blob: nil, err: expectedErr, } factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() tenantID := "tenant123" location := "file-abc.txt" blob, err := svc.DownloadAgentFile(tenantID, location) if err == nil { t.Fatalf("expected error, got nil") } if !errors.Is(err, expectedErr) { t.Errorf("expected error %v, got %v", expectedErr, err) } if blob != nil { t.Errorf("expected nil blob, got %v", blob) } } func TestFileService_UploadFromURL_PDFAddsExtensionAndStoresToDownloads(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/pdf") _, _ = w.Write([]byte("%PDF-1.7 fake pdf")) })) defer server.Close() origAssert := assertURLSafe origPinned := pinnedHTTPClient assertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } pinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { assertURLSafe = origAssert pinnedHTTPClient = origPinned }) mockStorage := &fakeStorage{} factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() resp, err := svc.UploadFromURL("tenant123", server.URL+"/report") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) } if mockStorage.lastBucket != "tenant123-downloads" { t.Fatalf("bucket = %q", mockStorage.lastBucket) } if resp["name"] != "report.pdf" { t.Fatalf("name = %#v, want report.pdf", resp["name"]) } if resp["mime_type"] != "application/pdf" { t.Fatalf("mime_type = %#v", resp["mime_type"]) } if resp["id"] != mockStorage.lastFnm { t.Fatalf("id = %#v, stored key = %q", resp["id"], mockStorage.lastFnm) } } func TestFileService_UploadFromURL_HTMLNormalizesReadableContent(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/html; charset=utf-8") _, _ = w.Write([]byte(`
Hello

World

`)) })) defer server.Close() origAssert := assertURLSafe origPinned := pinnedHTTPClient assertURLSafe = func(rawURL string) (string, string, error) { return "127.0.0.1", "127.0.0.1", nil } pinnedHTTPClient = func(hostname, resolvedIP string, timeout time.Duration) *http.Client { return server.Client() } t.Cleanup(func() { assertURLSafe = origAssert pinnedHTTPClient = origPinned }) mockStorage := &fakeStorage{} factory := storage.GetStorageFactory() originalStorage := factory.GetStorage() factory.SetStorage(mockStorage) t.Cleanup(func() { factory.SetStorage(originalStorage) }) svc := testFileService() resp, err := svc.UploadFromURL("tenant123", server.URL+"/page") if err != nil { t.Fatalf("UploadFromURL failed: %v", err) } stored := string(mockStorage.blob) if strings.Contains(strings.ToLower(stored), "