mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-10-12 02:16:39 +00:00
2314 lines
74 KiB
Go
2314 lines
74 KiB
Go
//
|
|
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
//
|
|
// Licensed under the Apache License, Version 2.0 (the "License");
|
|
// you may not use this file except in compliance with the License.
|
|
// You may obtain a copy of the License at
|
|
//
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
|
//
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
// See the License for the specific language governing permissions and
|
|
// limitations under the License.
|
|
//
|
|
|
|
package handler
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
"github.com/glebarez/sqlite"
|
|
"gorm.io/gorm"
|
|
|
|
"ragflow/internal/common"
|
|
"ragflow/internal/dao"
|
|
"ragflow/internal/entity"
|
|
"ragflow/internal/service"
|
|
dataset "ragflow/internal/service/dataset"
|
|
"ragflow/internal/service/document"
|
|
)
|
|
|
|
func TestParseMetadataSelectorDocumentIDs(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
body string
|
|
wantIDs []string
|
|
wantErr string
|
|
}{
|
|
{name: "missing selector", body: `{}`},
|
|
{name: "missing IDs", body: `{"selector":{}}`},
|
|
{name: "null IDs", body: `{"selector":{"document_ids":null}}`},
|
|
{name: "empty IDs", body: `{"selector":{"document_ids":[]}}`},
|
|
{name: "string IDs", body: `{"selector":{"document_ids":["doc-1","doc-2"]}}`, wantIDs: []string{"doc-1", "doc-2"}},
|
|
{name: "non-list IDs", body: `{"selector":{"document_ids":"doc-1"}}`, wantErr: "document_ids must be a list."},
|
|
{name: "numeric ID", body: `{"selector":{"document_ids":[123]}}`, wantErr: "document_ids must contain only strings."},
|
|
{name: "boolean ID", body: `{"selector":{"document_ids":[true]}}`, wantErr: "document_ids must contain only strings."},
|
|
{name: "null ID", body: `{"selector":{"document_ids":[null]}}`, wantErr: "document_ids must contain only strings."},
|
|
{name: "object ID", body: `{"selector":{"document_ids":[{}]}}`, wantErr: "document_ids must contain only strings."},
|
|
{name: "array ID", body: `{"selector":{"document_ids":[[]]}}`, wantErr: "document_ids must contain only strings."},
|
|
{name: "mixed IDs", body: `{"selector":{"document_ids":["doc-1",123]}}`, wantErr: "document_ids must contain only strings."},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
var body map[string]interface{}
|
|
if err := json.Unmarshal([]byte(tt.body), &body); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
selector, errMsg := parseMetadataSelector(body["selector"])
|
|
if errMsg != tt.wantErr {
|
|
t.Fatalf("error = %q, want %q", errMsg, tt.wantErr)
|
|
}
|
|
if tt.wantErr != "" {
|
|
if selector != nil {
|
|
t.Fatalf("selector = %#v, want nil", selector)
|
|
}
|
|
return
|
|
}
|
|
if selector == nil || !reflect.DeepEqual(selector.DocumentIDs, tt.wantIDs) {
|
|
t.Fatalf("selector = %#v, want IDs %v", selector, tt.wantIDs)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDocumentMetadataBatchRejectsNonStringIDs(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
h := &DocumentHandler{
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
tests := []struct {
|
|
method string
|
|
path string
|
|
handler gin.HandlerFunc
|
|
}{
|
|
{http.MethodPost, "/api/v1/datasets/kb-owner/metadata/update", h.MetadataBatchUpdate},
|
|
{http.MethodPatch, "/api/v1/datasets/kb-owner/documents/metadatas", h.UpdateDocumentMetadatas},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.method, func(t *testing.T) {
|
|
c, w := setupGinContextWithUser(tt.method, tt.path, `{"selector":{"document_ids":[123]}}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "kb-owner"}}
|
|
tt.handler(c)
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("status = %d, want 200: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp struct {
|
|
Code common.ErrorCode `json:"code"`
|
|
Message string `json:"message"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp.Code != common.CodeDataError || resp.Message != "document_ids must contain only strings." {
|
|
t.Fatalf("unexpected response: %s", w.Body.String())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// fakeDocumentService implements documentServiceIface for handler tests.
|
|
type fakeDocumentService struct {
|
|
deleted int
|
|
err error
|
|
downloadErr error
|
|
doc *document.DocumentResponse
|
|
docErr error
|
|
updateCalled bool
|
|
updatedID string
|
|
updateCode common.ErrorCode
|
|
updateErr error
|
|
deleteCalled bool
|
|
deletedID string
|
|
stopResult map[string]interface{}
|
|
stopErr error
|
|
stopIngestionTasks []*entity.IngestionTask
|
|
stopIngestionTaskErr error
|
|
removeIngestionTasks []map[string]string
|
|
removeIngestionTaskErr error
|
|
thumbnails map[string]string
|
|
thumbnailErr error
|
|
thumbnailUserID string
|
|
thumbnailDocIDs []string
|
|
metadataSummary map[string]interface{}
|
|
metadataErr error
|
|
metadataKBID string
|
|
metadataDocIDs []string
|
|
setMetaCalled bool
|
|
setMetaDocID string
|
|
setMetaValue map[string]interface{}
|
|
uploadLocalData []map[string]interface{}
|
|
uploadLocalErrs []string
|
|
uploadLocalKB *entity.Knowledgebase
|
|
uploadLocalPath string
|
|
uploadOverride map[string]interface{}
|
|
ingestCode common.ErrorCode
|
|
ingestErr error
|
|
ingestUserID string
|
|
ingestReq *document.IngestDocumentRequest
|
|
listOpts dao.DocumentListOptions
|
|
filterOpts dao.DocumentListOptions
|
|
filterResult map[string]interface{}
|
|
filterTotal int64
|
|
listIDs []string
|
|
documentList []*entity.DocumentListItem
|
|
documentListTotal int64
|
|
metadataByKBs map[string]interface{}
|
|
hasActiveTasks bool
|
|
hasActiveTasksErr error
|
|
imageData []byte
|
|
imageErr error
|
|
imageUserID string
|
|
imageID string
|
|
imageDocumentID string
|
|
}
|
|
|
|
func TestMapDocumentListItemIncludesLatestIngestionEvent(t *testing.T) {
|
|
event := &service.IngestionEventItem{ID: 42, EventType: dao.EventTypeMessage, Message: "queued"}
|
|
item := mapDocumentListItem(&entity.DocumentListItem{ID: "doc-1", KbID: "kb-1"}, map[string]interface{}{}, event)
|
|
got, ok := item["latest_ingestion_event"].(*service.IngestionEventItem)
|
|
if !ok || got != event {
|
|
t.Fatalf("latest_ingestion_event = %#v, want original event", item["latest_ingestion_event"])
|
|
}
|
|
}
|
|
|
|
func (f *fakeDocumentService) Ingest(ctx context.Context, userID string, req *document.IngestDocumentRequest) (common.ErrorCode, error) {
|
|
f.ingestUserID = userID
|
|
f.ingestReq = req
|
|
if f.ingestCode != 0 || f.ingestErr != nil {
|
|
return f.ingestCode, f.ingestErr
|
|
}
|
|
return common.CodeSuccess, nil
|
|
}
|
|
|
|
const uploadTestDatasetID = "123e4567-e89b-12d3-a456-426614174000"
|
|
|
|
func (f *fakeDocumentService) UpdateDatasetDocument(ctx context.Context, userID, datasetID, documentID string, req *document.UpdateDatasetDocumentRequest, present map[string]bool) (*document.UpdateDatasetDocumentResponse, common.ErrorCode, error) {
|
|
return nil, common.CodeSuccess, nil
|
|
}
|
|
func (f *fakeDocumentService) BatchUpdateDocumentMetadatas(ctx context.Context, datasetID string, selector *document.MetadataSelector, updates []document.MetadataUpdate, deletes []document.MetadataDelete) (*document.BatchUpdateMetadatasResponse, common.ErrorCode, error) {
|
|
return nil, common.CodeSuccess, nil
|
|
}
|
|
|
|
func (f *fakeDocumentService) GetDocumentArtifact(ctx context.Context, filename, _ string) (*document.ArtifactResponse, error) {
|
|
if filename == "error.txt" {
|
|
return nil, document.ErrArtifactNotFound
|
|
}
|
|
if filename == "unexpected.txt" {
|
|
return nil, fmt.Errorf("unexpected error")
|
|
}
|
|
return &document.ArtifactResponse{
|
|
Data: []byte("artifact content"),
|
|
ContentType: "text/plain",
|
|
SafeFilename: "safe.txt",
|
|
ForceAttachment: false,
|
|
}, nil
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentPreview(ctx context.Context, userID, docID string) (*document.DocumentPreview, error) {
|
|
if docID == "not-found" {
|
|
return nil, document.ErrPreviewDocumentNotFound
|
|
}
|
|
if docID == "empty-file" {
|
|
return nil, document.ErrPreviewFileEmpty
|
|
}
|
|
if docID == "storage-error" {
|
|
return nil, fmt.Errorf("read document object b/k: connection refused")
|
|
}
|
|
return &document.DocumentPreview{
|
|
Data: []byte("preview content"),
|
|
ContentType: "text/plain",
|
|
FileName: "preview.txt",
|
|
}, nil
|
|
}
|
|
func (f *fakeDocumentService) DownloadDocument(ctx context.Context, datasetID, docID string) (*document.DownloadDocumentResp, error) {
|
|
if f.downloadErr != nil {
|
|
return nil, f.downloadErr
|
|
}
|
|
if docID == "not-found" {
|
|
return nil, document.ErrDocumentNotFound
|
|
}
|
|
return &document.DownloadDocumentResp{
|
|
Data: []byte("document data"),
|
|
ContentType: "application/pdf",
|
|
FileName: "doc.pdf",
|
|
}, nil
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentByID(ctx context.Context, id string) (*document.DocumentResponse, error) {
|
|
if f.docErr != nil {
|
|
return nil, f.docErr
|
|
}
|
|
if f.doc != nil {
|
|
return f.doc, nil
|
|
}
|
|
return nil, fmt.Errorf("document not found")
|
|
}
|
|
func (f *fakeDocumentService) UpdateDocument(ctx context.Context, id string, req *document.UpdateDocumentRequest) (common.ErrorCode, error) {
|
|
f.updateCalled = true
|
|
f.updatedID = id
|
|
return f.updateCode, f.updateErr
|
|
}
|
|
func (f *fakeDocumentService) DeleteDocument(ctx context.Context, id string) error {
|
|
f.deleteCalled = true
|
|
f.deletedID = id
|
|
return nil
|
|
}
|
|
func (f *fakeDocumentService) DeleteDocuments(ctx context.Context, ids []string, deleteAll bool, datasetID, userID string) (int, error) {
|
|
return f.deleted, f.err
|
|
}
|
|
func (f *fakeDocumentService) ParseDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeDocumentService) StopParseDocuments(ctx context.Context, datasetID string, docIDs []string) (map[string]interface{}, error) {
|
|
return f.stopResult, f.stopErr
|
|
}
|
|
func (f *fakeDocumentService) ListDocuments(ctx context.Context, page, pageSize int) ([]*document.DocumentResponse, int64, error) {
|
|
return nil, 0, nil
|
|
}
|
|
func (f *fakeDocumentService) ListDocumentsByDatasetID(ctx context.Context, kbID, keywords string, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
|
|
return nil, 0, nil
|
|
}
|
|
func (f *fakeDocumentService) ListDocumentsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions, page, pageSize int) ([]*entity.DocumentListItem, int64, error) {
|
|
f.listOpts = opts
|
|
return f.documentList, f.documentListTotal, nil
|
|
}
|
|
func (f *fakeDocumentService) ListDocumentIDsByDatasetIDWithOptions(ctx context.Context, opts dao.DocumentListOptions) ([]string, error) {
|
|
f.listOpts = opts
|
|
return f.listIDs, nil
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentFiltersByDatasetID(ctx context.Context, opts dao.DocumentListOptions) (map[string]interface{}, int64, error) {
|
|
f.filterOpts = opts
|
|
if f.filterResult != nil {
|
|
return f.filterResult, f.filterTotal, nil
|
|
}
|
|
return map[string]interface{}{}, 0, nil
|
|
}
|
|
func (f *fakeDocumentService) GetMetadataByKBs(ctx context.Context, kbIDs []string) (map[string]interface{}, error) {
|
|
if f.metadataByKBs != nil {
|
|
return f.metadataByKBs, nil
|
|
}
|
|
return map[string]interface{}{}, nil
|
|
}
|
|
func (f *fakeDocumentService) BatchUpdateDocumentStatus(ctx context.Context, userID, datasetID, status string, documentIDs []string) (map[string]interface{}, common.ErrorCode, error) {
|
|
return map[string]interface{}{}, common.CodeSuccess, nil
|
|
}
|
|
func (f *fakeDocumentService) GetThumbnails(ctx context.Context, userID string, docIDs []string) (map[string]string, error) {
|
|
f.thumbnailUserID = userID
|
|
f.thumbnailDocIDs = append([]string(nil), docIDs...)
|
|
return f.thumbnails, f.thumbnailErr
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentImage(ctx context.Context, userID, imageID string) ([]byte, error) {
|
|
f.imageUserID = userID
|
|
f.imageID = imageID
|
|
return f.imageData, f.imageErr
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentImageForDocument(ctx context.Context, userID, docID, imageID string) ([]byte, error) {
|
|
f.imageUserID = userID
|
|
f.imageDocumentID = docID
|
|
f.imageID = imageID
|
|
return f.imageData, f.imageErr
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentThumbnail(ctx context.Context, userID, docID string) ([]byte, error) {
|
|
f.imageUserID = userID
|
|
f.imageDocumentID = docID
|
|
return f.imageData, f.imageErr
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentsByAuthorID(ctx context.Context, authorID, page, pageSize int) ([]*document.DocumentResponse, int64, error) {
|
|
return nil, 0, nil
|
|
}
|
|
func (f *fakeDocumentService) GetMetadataSummary(ctx context.Context, kbID string, docIDs []string) (map[string]interface{}, error) {
|
|
f.metadataKBID = kbID
|
|
f.metadataDocIDs = docIDs
|
|
return f.metadataSummary, f.metadataErr
|
|
}
|
|
func (f *fakeDocumentService) HasActiveIngestionTasks(ctx context.Context, datasetID string) (bool, error) {
|
|
return f.hasActiveTasks, f.hasActiveTasksErr
|
|
}
|
|
|
|
func (f *fakeDocumentService) SetDocumentMetadata(ctx context.Context, docID string, meta map[string]interface{}) error {
|
|
f.setMetaCalled = true
|
|
f.setMetaDocID = docID
|
|
f.setMetaValue = meta
|
|
return nil
|
|
}
|
|
func (f *fakeDocumentService) DeleteDocumentMetadata(ctx context.Context, docID string, keys []string) error {
|
|
return nil
|
|
}
|
|
func (f *fakeDocumentService) DeleteDocumentAllMetadata(ctx context.Context, docID string) error {
|
|
return nil
|
|
}
|
|
func (f *fakeDocumentService) GetDocumentMetadataByID(ctx context.Context, docID string) (map[string]interface{}, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeDocumentService) UploadLocalDocuments(ctx context.Context, kb *entity.Knowledgebase, tenantID string, files []*multipart.FileHeader, parentPath string, parserConfigOverride map[string]interface{}) ([]map[string]interface{}, []string) {
|
|
f.uploadLocalKB = kb
|
|
f.uploadLocalPath = parentPath
|
|
f.uploadOverride = parserConfigOverride
|
|
return f.uploadLocalData, f.uploadLocalErrs
|
|
}
|
|
func (f *fakeDocumentService) UploadWebDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name, url string) (map[string]interface{}, common.ErrorCode, error) {
|
|
return nil, common.CodeServerError, fmt.Errorf("not implemented")
|
|
}
|
|
func (f *fakeDocumentService) UploadEmptyDocument(ctx context.Context, kb *entity.Knowledgebase, tenantID, name string) (map[string]interface{}, common.ErrorCode, error) {
|
|
return nil, common.CodeServerError, fmt.Errorf("not implemented")
|
|
}
|
|
|
|
func (f *fakeDocumentService) ListIngestionTasks(ctx context.Context, userID string, datasetID *string, page, pageSize int) ([]*entity.IngestionTask, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeDocumentService) IngestDocuments(ctx context.Context, datasetID, userID string, docIDs []string) ([]*service.ParseDocumentResponse, error) {
|
|
return nil, nil
|
|
}
|
|
func (f *fakeDocumentService) StopIngestionTasks(ctx context.Context, tasks []string, userID string) ([]*entity.IngestionTask, error) {
|
|
return f.stopIngestionTasks, f.stopIngestionTaskErr
|
|
}
|
|
func (f *fakeDocumentService) RemoveIngestionTasks(ctx context.Context, tasks []string, userID string) ([]map[string]string, error) {
|
|
return f.removeIngestionTasks, f.removeIngestionTaskErr
|
|
}
|
|
|
|
func setupGinContextWithUser(method, path, body string) (*gin.Context, *httptest.ResponseRecorder) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = req
|
|
c.Set("user", &entity.User{ID: "user-1"})
|
|
c.Set("user_id", "user-1")
|
|
return c, w
|
|
}
|
|
|
|
func setupDocumentPermissionDB(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.Knowledgebase{},
|
|
&entity.UserTenant{},
|
|
); err != nil {
|
|
t.Fatalf("failed to migrate permission tables: %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.UserTenant{
|
|
ID: "ut-owner",
|
|
UserID: "owner-user",
|
|
TenantID: "tenant-owner",
|
|
Role: "owner",
|
|
Status: sptr(string(entity.StatusValid)),
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert tenant owner: %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 TestDocumentReadHandlersForbidden(t *testing.T) {
|
|
setupDocumentPermissionDB(t, false)
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
tests := []struct {
|
|
name string
|
|
method string
|
|
path string
|
|
body string
|
|
params gin.Params
|
|
handler gin.HandlerFunc
|
|
}{
|
|
{"document detail", http.MethodGet, "/api/v1/documents/doc-1", "", gin.Params{{Key: "id", Value: "doc-1"}}, h.GetDocumentByID},
|
|
{"metadata summary", http.MethodPost, "/api/v1/document/metadata/summary", `{"kb_id":"kb-owner","doc_ids":["doc-1"]}`, nil, h.MetadataSummary},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
c, w := setupGinContextWithUser(tt.method, tt.path, tt.body)
|
|
c.Params = tt.params
|
|
tt.handler(c)
|
|
if w.Code != http.StatusForbidden {
|
|
t.Fatalf("status = %d, want 403: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp struct {
|
|
Code common.ErrorCode `json:"code"`
|
|
Message string `json:"message"`
|
|
Data interface{} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp.Code != common.CodeForbidden || resp.Message != "Permission denied" || resp.Data != nil {
|
|
t.Fatalf("unexpected response: %s", w.Body.String())
|
|
}
|
|
if fake.metadataKBID != "" {
|
|
t.Fatal("GetMetadataSummary should not be called without dataset access")
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestSetMetaHandler_NotAccessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, false)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/document/set_meta", `{"doc_id":"doc-1","meta":"{\"poc\":\"blocked\"}"}`)
|
|
h.SetMeta(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeForbidden) {
|
|
t.Fatalf("expected auth error, got %v", resp)
|
|
}
|
|
if resp["message"] != "Permission denied" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
if fake.setMetaCalled {
|
|
t.Fatal("SetDocumentMetadata should not be called without dataset access")
|
|
}
|
|
}
|
|
|
|
func TestSetMetaHandler_Accessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/document/set_meta", `{"doc_id":"doc-1","meta":"{\"category\":\"tech\",\"year\":2026}"}`)
|
|
h.SetMeta(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) || resp["data"] != true {
|
|
t.Fatalf("unexpected response: %v", resp)
|
|
}
|
|
if !fake.setMetaCalled {
|
|
t.Fatal("SetDocumentMetadata should be called with dataset access")
|
|
}
|
|
if fake.setMetaDocID != "doc-1" {
|
|
t.Fatalf("set meta doc id = %q, want doc-1", fake.setMetaDocID)
|
|
}
|
|
if fake.setMetaValue["category"] != "tech" || fake.setMetaValue["year"] != float64(2026) {
|
|
t.Fatalf("unexpected meta: %#v", fake.setMetaValue)
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentHandler_NotAccessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, false)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/documents/doc-1", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.DeleteDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeForbidden) {
|
|
t.Fatalf("expected auth error, got %v", resp)
|
|
}
|
|
if resp["message"] != "Permission denied" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
if fake.deleteCalled {
|
|
t.Fatal("DeleteDocument should not be called without dataset access")
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentHandler_Accessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/documents/doc-1", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.DeleteDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if !fake.deleteCalled {
|
|
t.Fatal("DeleteDocument should be called with dataset access")
|
|
}
|
|
if fake.deletedID != "doc-1" {
|
|
t.Fatalf("deleted id = %q, want doc-1", fake.deletedID)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["message"] != "deleted successfully" {
|
|
t.Fatalf("unexpected response: %v", resp)
|
|
}
|
|
}
|
|
|
|
func TestUpdateDocumentHandler_NotAccessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, false)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"name":"blocked"}`)
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.UpdateDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeForbidden) {
|
|
t.Fatalf("expected auth error, got %v", resp)
|
|
}
|
|
if resp["message"] != "Permission denied" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
if fake.updateCalled {
|
|
t.Fatal("UpdateDocument should not be called without dataset access")
|
|
}
|
|
}
|
|
|
|
func TestUpdateDocumentHandler_Accessible(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"name":"allowed"}`)
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.UpdateDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if !fake.updateCalled {
|
|
t.Fatal("UpdateDocument should be called with dataset access")
|
|
}
|
|
if fake.updatedID != "doc-1" {
|
|
t.Fatalf("updated id = %q, want doc-1", fake.updatedID)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["message"] != "updated successfully" {
|
|
t.Fatalf("unexpected response: %v", resp)
|
|
}
|
|
}
|
|
|
|
func TestUpdateDocumentHandler_ValidationError(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
updateCode: common.CodeDataError,
|
|
updateErr: errors.New("can't change `progress`"),
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"progress":0.5}`)
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.UpdateDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeDataError) {
|
|
t.Fatalf("expected data error, got %v", resp)
|
|
}
|
|
if resp["message"] != "can't change `progress`" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestUpdateDocumentHandler_ServerError(t *testing.T) {
|
|
setupDocumentPermissionDB(t, true)
|
|
|
|
fake := &fakeDocumentService{
|
|
doc: &document.DocumentResponse{ID: "doc-1", KbID: "kb-owner"},
|
|
updateCode: common.CodeServerError,
|
|
updateErr: errors.New("database unavailable"),
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/documents/doc-1", `{"name":"new.pdf"}`)
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
h.UpdateDocument(c)
|
|
|
|
if w.Code != http.StatusInternalServerError {
|
|
t.Fatalf("expected 500, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeServerError) {
|
|
t.Fatalf("expected server error, got %v", resp)
|
|
}
|
|
if resp["message"] != "database unavailable" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
func setupUploadHandlerDB(t *testing.T, role string) *gorm.DB {
|
|
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.User{},
|
|
&entity.Tenant{},
|
|
&entity.UserTenant{},
|
|
&entity.Knowledgebase{},
|
|
); err != nil {
|
|
t.Fatalf("failed to migrate: %v", err)
|
|
}
|
|
if err := db.Create(&entity.User{ID: "user-1", Nickname: "test", Email: "[email protected]", Password: sptr("x")}).Error; err != nil {
|
|
t.Fatalf("insert user: %v", err)
|
|
}
|
|
if err := db.Create(&entity.Tenant{ID: "tenant-1", LLMID: "llm-1", EmbdID: "embd-1", ASRID: "asr-1", Status: sptr(string(entity.StatusValid))}).Error; err != nil {
|
|
t.Fatalf("insert tenant: %v", err)
|
|
}
|
|
if err := db.Create(&entity.UserTenant{ID: "ut-1", UserID: "user-1", TenantID: "tenant-1", Role: role, Status: sptr(string(entity.StatusValid))}).Error; err != nil {
|
|
t.Fatalf("insert user_tenant: %v", err)
|
|
}
|
|
if err := db.Create(&entity.UserTenant{ID: "ut-owner", UserID: "owner-1", TenantID: "tenant-1", Role: "owner", Status: sptr(string(entity.StatusValid))}).Error; err != nil {
|
|
t.Fatalf("insert tenant owner: %v", err)
|
|
}
|
|
pipelineID := "pipe-1"
|
|
if err := db.Create(&entity.Knowledgebase{
|
|
ID: "123e4567e89b12d3a456426614174000",
|
|
TenantID: "tenant-1",
|
|
Name: "kb-upload",
|
|
EmbdID: "embd-1",
|
|
CreatedBy: "user-1",
|
|
Permission: string(entity.TenantPermissionTeam),
|
|
ParserID: "naive",
|
|
PipelineID: &pipelineID,
|
|
ParserConfig: entity.JSONMap{"base": "cfg"},
|
|
Status: sptr(string(entity.StatusValid)),
|
|
}).Error; err != nil {
|
|
t.Fatalf("insert knowledgebase: %v", err)
|
|
}
|
|
return db
|
|
}
|
|
|
|
func setupUploadContext(t *testing.T, path string, fields map[string]string, fileName string, fileContent []byte) (*gin.Context, *httptest.ResponseRecorder) {
|
|
t.Helper()
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
var body bytes.Buffer
|
|
writer := multipart.NewWriter(&body)
|
|
for k, v := range fields {
|
|
if err := writer.WriteField(k, v); err != nil {
|
|
t.Fatalf("write field %s: %v", k, err)
|
|
}
|
|
}
|
|
part, err := writer.CreateFormFile("file", fileName)
|
|
if err != nil {
|
|
t.Fatalf("create form file: %v", err)
|
|
}
|
|
if _, err := part.Write(fileContent); err != nil {
|
|
t.Fatalf("write form file: %v", err)
|
|
}
|
|
if err := writer.Close(); err != nil {
|
|
t.Fatalf("close writer: %v", err)
|
|
}
|
|
req := httptest.NewRequest(http.MethodPost, path, &body)
|
|
req.Header.Set("Content-Type", writer.FormDataContentType())
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = req
|
|
c.Set("user", &entity.User{ID: "user-1"})
|
|
c.Set("user_id", "user-1")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: uploadTestDatasetID}}
|
|
return c, w
|
|
}
|
|
|
|
func setupDocumentIngestRoute(userID string, svc *fakeDocumentService) *gin.Engine {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: svc,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
r := gin.New()
|
|
r.Use(func(c *gin.Context) {
|
|
c.Set("user", &entity.User{ID: userID})
|
|
c.Set("user_id", userID)
|
|
})
|
|
r.POST("/api/v1/documents/ingest", h.Ingest)
|
|
return r
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_Success(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{deleted: 3}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/documents", `{"ids": ["doc-1", "doc-2", "doc-3"]}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected code 0, got %v", resp["code"])
|
|
}
|
|
data := resp["data"].(map[string]interface{})
|
|
if data["deleted"] != float64(3) {
|
|
t.Fatalf("expected deleted=3, got %v", data["deleted"])
|
|
}
|
|
}
|
|
|
|
func TestUploadDocumentsHandler_LocalUsesFullKBAndIgnoresBadParserConfig(t *testing.T) {
|
|
db := setupUploadHandlerDB(t, "normal")
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
uploadLocalData: []map[string]interface{}{
|
|
{"id": "doc-1", "kb_id": "ds-1", "parser_id": "naive", "chunk_num": int64(0), "token_num": int64(0), "name": "a.txt"},
|
|
},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupUploadContext(t, "/api/v1/datasets/ds-1/documents?type=local", map[string]string{
|
|
"parent_path": "nested/path",
|
|
"parser_config": "{bad json",
|
|
}, "a.txt", []byte("abc"))
|
|
|
|
h.UploadDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if fake.uploadLocalKB == nil {
|
|
t.Fatalf("UploadLocalDocuments was not called, response=%s", w.Body.String())
|
|
}
|
|
if fake.uploadLocalKB.TenantID != "tenant-1" || fake.uploadLocalKB.Name != "kb-upload" || fake.uploadLocalKB.ParserID != "naive" {
|
|
t.Fatalf("incomplete kb passed to service: %+v", fake.uploadLocalKB)
|
|
}
|
|
if fake.uploadLocalPath != "nested/path" {
|
|
t.Fatalf("parent path=%q, want nested/path", fake.uploadLocalPath)
|
|
}
|
|
if fake.uploadOverride != nil {
|
|
t.Fatalf("bad parser_config should be ignored, got %v", fake.uploadOverride)
|
|
}
|
|
}
|
|
|
|
func TestUploadDocumentsHandler_LocalReturnsPartialSuccess(t *testing.T) {
|
|
db := setupUploadHandlerDB(t, "normal")
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
uploadLocalData: []map[string]interface{}{
|
|
{"id": "doc-1", "kb_id": "ds-1", "parser_id": "naive", "chunk_num": int64(0), "token_num": int64(0), "name": "ok.txt"},
|
|
},
|
|
uploadLocalErrs: []string{"bad.exe: This type of file has not been supported yet!"},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupUploadContext(t, "/api/v1/datasets/ds-1/documents?type=local", nil, "ok.txt", []byte("abc"))
|
|
h.UploadDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if resp["code"] != float64(common.CodeServerError) {
|
|
t.Fatalf("expected server error code for partial upload, got %v", resp)
|
|
}
|
|
data := resp["data"].([]interface{})
|
|
if len(data) != 1 {
|
|
t.Fatalf("expected one successful document, got %v", data)
|
|
}
|
|
if !strings.Contains(resp["message"].(string), "bad.exe") {
|
|
t.Fatalf("expected failed file in message, got %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestUploadDocumentsHandler_DeniesNonNormalTeamRole(t *testing.T) {
|
|
db := setupUploadHandlerDB(t, "admin")
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupUploadContext(t, "/api/v1/datasets/ds-1/documents?type=local", nil, "a.txt", []byte("abc"))
|
|
h.UploadDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if resp["code"] == float64(common.CodeSuccess) {
|
|
t.Fatalf("expected authorization error, got %v", resp)
|
|
}
|
|
if fake.uploadLocalKB != nil {
|
|
t.Fatal("service should not be called on denied upload")
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_DeleteAll(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{deleted: 5}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/documents", `{"delete_all": true}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_MutuallyExclusive(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/documents", `{"ids": ["doc-1"], "delete_all": true}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for mutually exclusive ids+delete_all")
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_NoIDsNoDeleteAll(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/documents", `{}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for no ids and no delete_all")
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_ServiceError(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{err: fmt.Errorf("permission denied")}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/documents", `{"ids": ["doc-1"]}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error code")
|
|
}
|
|
}
|
|
|
|
func TestDeleteDocumentsHandler_MissingDatasetID(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets//documents", `{"ids": ["doc-1"]}`)
|
|
|
|
h.DeleteDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for missing dataset_id")
|
|
}
|
|
}
|
|
|
|
func TestDocumentHandlerIngestMatchesPythonResponseShape(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/documents/ingest", `{"doc_ids":["doc-1"],"run":"1"}`)
|
|
h.Ingest(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected top-level code 0, got %v", resp["code"])
|
|
}
|
|
if resp["data"] != true {
|
|
t.Fatalf("expected top-level data=true, got %#v", resp["data"])
|
|
}
|
|
if _, ok := resp["data"].(map[string]interface{}); ok {
|
|
t.Fatalf("response must not nest code/message under data: %#v", resp["data"])
|
|
}
|
|
if fake.ingestUserID != "user-1" {
|
|
t.Fatalf("expected user-1, got %q", fake.ingestUserID)
|
|
}
|
|
if fake.ingestReq == nil || len(fake.ingestReq.DocIDs) != 1 || fake.ingestReq.DocIDs[0] != "doc-1" {
|
|
t.Fatalf("unexpected ingest request: %#v", fake.ingestReq)
|
|
}
|
|
}
|
|
|
|
func TestDocumentIngestRoutePassesPythonBodyToService(t *testing.T) {
|
|
fake := &fakeDocumentService{}
|
|
r := setupDocumentIngestRoute("user-1", fake)
|
|
|
|
w := httptest.NewRecorder()
|
|
req := httptest.NewRequest("POST", "/api/v1/documents/ingest", strings.NewReader(`{"doc_ids":["doc-1","doc-2"],"run":1,"delete":true,"apply_kb":true}`))
|
|
req.Header.Set("Content-Type", "application/json")
|
|
r.ServeHTTP(w, req)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) || resp["data"] != true {
|
|
t.Fatalf("unexpected response: %s", w.Body.String())
|
|
}
|
|
if fake.ingestUserID != "user-1" {
|
|
t.Fatalf("userID = %q, want user-1", fake.ingestUserID)
|
|
}
|
|
if fake.ingestReq == nil {
|
|
t.Fatal("service did not receive ingest request")
|
|
}
|
|
if len(fake.ingestReq.DocIDs) != 2 || fake.ingestReq.DocIDs[0] != "doc-1" || fake.ingestReq.DocIDs[1] != "doc-2" {
|
|
t.Fatalf("doc_ids = %#v, want [doc-1 doc-2]", fake.ingestReq.DocIDs)
|
|
}
|
|
if fmt.Sprint(fake.ingestReq.Run) != "1" {
|
|
t.Fatalf("run = %#v, want 1", fake.ingestReq.Run)
|
|
}
|
|
if !fake.ingestReq.Delete {
|
|
t.Fatal("delete = false, want true")
|
|
}
|
|
if !fake.ingestReq.ApplyKB {
|
|
t.Fatal("apply_kb = false, want true")
|
|
}
|
|
}
|
|
|
|
func TestDocumentHandlerIngestPropagatesServiceErrorCode(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
ingestCode: common.CodeAuthenticationError,
|
|
ingestErr: fmt.Errorf("no authorization"),
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/documents/ingest", `{"doc_ids":["doc-1"],"run":"1"}`)
|
|
h.Ingest(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeAuthenticationError) {
|
|
t.Fatalf("expected auth error code, got %v", resp["code"])
|
|
}
|
|
if resp["message"] != "no authorization" {
|
|
t.Fatalf("unexpected message: %v", resp["message"])
|
|
}
|
|
if resp["data"] != nil {
|
|
t.Fatalf("expected nil data, got %#v", resp["data"])
|
|
}
|
|
}
|
|
|
|
func TestDocumentHandlerIngest_CancelDispatchesToService(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/documents/ingest", `{"doc_ids":["doc-1"],"run":2}`)
|
|
h.Ingest(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected code 0, got %v", resp["code"])
|
|
}
|
|
if resp["data"] != true {
|
|
t.Fatalf("expected data true, got %#v", resp["data"])
|
|
}
|
|
if fake.ingestReq == nil || len(fake.ingestReq.DocIDs) != 1 || fake.ingestReq.DocIDs[0] != "doc-1" {
|
|
t.Fatalf("unexpected ingestReq: %#v", fake.ingestReq)
|
|
}
|
|
if fmt.Sprint(fake.ingestReq.Run) != "2" {
|
|
t.Fatalf("run = %v, want 2", fake.ingestReq.Run)
|
|
}
|
|
}
|
|
|
|
func TestStopParseDocumentsHandler_EmptyDocIDs(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": []}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.StopParseDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for empty document_ids")
|
|
}
|
|
}
|
|
|
|
func TestStopParseDocumentsHandler_BadJSON(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `not json`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.StopParseDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for bad JSON body")
|
|
}
|
|
}
|
|
|
|
// setupHandlerAccessDB sets up SQLite in-memory DB for handler dataset access tests.
|
|
func setupHandlerAccessDB(t *testing.T) *gorm.DB {
|
|
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.User{},
|
|
&entity.Tenant{},
|
|
&entity.UserTenant{},
|
|
&entity.Knowledgebase{},
|
|
); err != nil {
|
|
t.Fatalf("failed to migrate: %v", err)
|
|
}
|
|
|
|
// Insert user
|
|
db.Create(&entity.User{ID: "user-1", Nickname: "test", Email: "[email protected]", Password: sptr("x")})
|
|
// Insert tenant
|
|
db.Create(&entity.Tenant{ID: "tenant-1", LLMID: "llm-1", EmbdID: "embd-1", ASRID: "asr-1"})
|
|
// Insert user_tenant mapping
|
|
db.Create(&entity.UserTenant{ID: "ut-owner", UserID: "owner-1", TenantID: "tenant-1", Role: "owner", Status: sptr(string(entity.StatusValid))})
|
|
db.Create(&entity.UserTenant{ID: "ut-1", UserID: "user-1", TenantID: "tenant-1", Role: "normal", Status: sptr(string(entity.StatusValid))})
|
|
// Insert knowledgebase
|
|
db.Create(&entity.Knowledgebase{
|
|
ID: "ds-1", TenantID: "tenant-1", Name: "test-kb", EmbdID: "embd-1",
|
|
CreatedBy: "user-1", Permission: string(entity.TenantPermissionTeam),
|
|
Status: sptr(string(entity.StatusValid)),
|
|
})
|
|
|
|
return db
|
|
}
|
|
|
|
// sptr returns a pointer to the given string (copy of service test helper).
|
|
func sptr(s string) *string { return &s }
|
|
|
|
func TestListDocumentsHandler_FilterRequestUsesQueryFilters(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
filterResult: map[string]interface{}{
|
|
"suffix": map[string]int64{"pdf": 2},
|
|
"run_status": map[string]int64{"3": 2},
|
|
"metadata": map[string]interface{}{},
|
|
},
|
|
filterTotal: 2,
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents?type=filter&keywords=report&suffix=pdf&run=COMPLETED&types=doc&desc=false", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if fake.filterOpts.KbID != "ds-1" {
|
|
t.Fatalf("expected dataset filter ds-1, got %q", fake.filterOpts.KbID)
|
|
}
|
|
if fake.filterOpts.Keywords != "report" {
|
|
t.Fatalf("expected keywords report, got %q", fake.filterOpts.Keywords)
|
|
}
|
|
if len(fake.filterOpts.Suffixes) != 1 || fake.filterOpts.Suffixes[0] != "pdf" {
|
|
t.Fatalf("expected suffix pdf, got %#v", fake.filterOpts.Suffixes)
|
|
}
|
|
if len(fake.filterOpts.RunStatuses) != 1 || fake.filterOpts.RunStatuses[0] != common.COMPLETED {
|
|
t.Fatalf("expected run COMPLETED to map to %q, got %#v", common.COMPLETED, fake.filterOpts.RunStatuses)
|
|
}
|
|
if len(fake.filterOpts.Types) != 1 || fake.filterOpts.Types[0] != "doc" {
|
|
t.Fatalf("expected type doc, got %#v", fake.filterOpts.Types)
|
|
}
|
|
if fake.filterOpts.Desc {
|
|
t.Fatal("expected desc=false to be parsed")
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("invalid json response: %v", err)
|
|
}
|
|
data := resp["data"].(map[string]interface{})
|
|
if data["total"] != float64(2) {
|
|
t.Fatalf("expected total 2, got %v", data["total"])
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsRejectsInvalidRunFilter(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
for _, invalid := range []string{"1", "DONE", "CANCEL", "FAIL", "SCHEDULE"} {
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents?run="+invalid, "user-1")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200 wrapper, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("invalid json response: %v", err)
|
|
}
|
|
if int(resp["code"].(float64)) != int(common.CodeDataError) {
|
|
t.Fatalf("expected code %d for %s, got %v", common.CodeDataError, invalid, resp["code"])
|
|
}
|
|
msg, _ := resp["message"].(string)
|
|
if !strings.Contains(msg, "Invalid filter run status conditions: "+invalid) {
|
|
t.Fatalf("expected error message for %s to contain 'Invalid filter run status conditions: %s', got %q", invalid, invalid, msg)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandlerReturnsScheduledIngestionStatus(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
documentList: []*entity.DocumentListItem{{
|
|
ID: "doc-scheduled",
|
|
KbID: "ds-1",
|
|
ParserID: "naive",
|
|
ParserConfig: "{}",
|
|
SourceType: "local",
|
|
Type: "document",
|
|
CreatedBy: "user-1",
|
|
Name: sptr("scheduled.pdf"),
|
|
Suffix: "pdf",
|
|
IngestionStatus: sptr(common.SCHEDULED),
|
|
}},
|
|
documentListTotal: 1,
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var response struct {
|
|
Data struct {
|
|
Docs []map[string]interface{} `json:"docs"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if len(response.Data.Docs) != 1 {
|
|
t.Fatalf("document count = %d, want 1", len(response.Data.Docs))
|
|
}
|
|
if got := response.Data.Docs[0]["ingestion_status"]; got != common.SCHEDULED {
|
|
t.Fatalf("ingestion_status = %v, want %q", got, common.SCHEDULED)
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandlerReturnsUnstartWhenNoIngestionTask(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
documentList: []*entity.DocumentListItem{{
|
|
ID: "doc-without-task",
|
|
KbID: "ds-1",
|
|
ParserID: "naive",
|
|
ParserConfig: "{}",
|
|
SourceType: "local",
|
|
Type: "document",
|
|
CreatedBy: "user-1",
|
|
Name: sptr("without-task.pdf"),
|
|
Suffix: "pdf",
|
|
}},
|
|
documentListTotal: 1,
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var response struct {
|
|
Data struct {
|
|
Docs []map[string]interface{} `json:"docs"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if len(response.Data.Docs) != 1 {
|
|
t.Fatalf("document count = %d, want 1", len(response.Data.Docs))
|
|
}
|
|
if got := response.Data.Docs[0]["ingestion_status"]; got != "UNSTART" {
|
|
t.Fatalf("ingestion_status = %v, want UNSTART", got)
|
|
}
|
|
if _, exists := response.Data.Docs[0]["run"]; exists {
|
|
t.Fatalf("run field should not exist in response")
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandlerReturnsHasActiveTasks(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
documentList: []*entity.DocumentListItem{{
|
|
ID: "doc-1",
|
|
KbID: "ds-1",
|
|
ParserID: "naive",
|
|
ParserConfig: "{}",
|
|
SourceType: "local",
|
|
Type: "document",
|
|
CreatedBy: "user-1",
|
|
Name: sptr("doc.pdf"),
|
|
Suffix: "pdf",
|
|
}},
|
|
documentListTotal: 1,
|
|
hasActiveTasks: true,
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var response struct {
|
|
Data struct {
|
|
Docs []map[string]interface{} `json:"docs"`
|
|
HasActiveTasks bool `json:"has_active_tasks"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if !response.Data.HasActiveTasks {
|
|
t.Fatalf("has_active_tasks = false, want true")
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandlerReturnsHasActiveTasksFalse(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
documentList: []*entity.DocumentListItem{{
|
|
ID: "doc-1",
|
|
KbID: "ds-1",
|
|
ParserID: "naive",
|
|
ParserConfig: "{}",
|
|
SourceType: "local",
|
|
Type: "document",
|
|
CreatedBy: "user-1",
|
|
Name: sptr("doc.pdf"),
|
|
Suffix: "pdf",
|
|
}},
|
|
documentListTotal: 1,
|
|
hasActiveTasks: false,
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var response struct {
|
|
Data struct {
|
|
Docs []map[string]interface{} `json:"docs"`
|
|
HasActiveTasks bool `json:"has_active_tasks"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if response.Data.HasActiveTasks {
|
|
t.Fatalf("has_active_tasks = true, want false")
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandlerKeepsPollingWhenActiveTaskCheckFails(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
fake := &fakeDocumentService{
|
|
documentList: []*entity.DocumentListItem{{
|
|
ID: "doc-1",
|
|
KbID: "ds-1",
|
|
ParserID: "naive",
|
|
ParserConfig: "{}",
|
|
SourceType: "local",
|
|
Type: "document",
|
|
CreatedBy: "user-1",
|
|
Name: sptr("doc.pdf"),
|
|
Suffix: "pdf",
|
|
}},
|
|
documentListTotal: 1,
|
|
hasActiveTasksErr: errors.New("temporary task lookup failure"),
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
var response struct {
|
|
Data struct {
|
|
HasActiveTasks bool `json:"has_active_tasks"`
|
|
} `json:"data"`
|
|
}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatalf("unmarshal response: %v", err)
|
|
}
|
|
if !response.Data.HasActiveTasks {
|
|
t.Fatal("has_active_tasks = false, want true when active task lookup fails")
|
|
}
|
|
}
|
|
|
|
func TestListDocumentsHandler_MetadataFilterNarrowsDocumentIDs(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
listIDs: []string{"doc-1", "doc-2", "doc-3"},
|
|
metadataByKBs: map[string]interface{}{
|
|
"author": map[string][]string{
|
|
"Alice": {"doc-2", "doc-4"},
|
|
},
|
|
},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/documents?metadata[author][]=Alice", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.ListDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if !fake.listOpts.DocIDFilterApplied {
|
|
t.Fatal("expected metadata filter to apply doc id filter")
|
|
}
|
|
if len(fake.listOpts.DocIDs) != 1 || fake.listOpts.DocIDs[0] != "doc-2" {
|
|
t.Fatalf("expected metadata filter to keep doc-2, got %#v", fake.listOpts.DocIDs)
|
|
}
|
|
}
|
|
|
|
func TestStopParseDocumentsHandler_Success(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
stopResult: map[string]interface{}{"success_count": 1},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.StopParseDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected code 0, got %v: %v", resp["code"], resp)
|
|
}
|
|
data := resp["data"].(map[string]interface{})
|
|
if data["success_count"] != float64(1) {
|
|
t.Fatalf("expected success_count=1, got %v", data["success_count"])
|
|
}
|
|
}
|
|
|
|
func TestStopParseDocumentsHandler_ServiceError(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
stopErr: fmt.Errorf("internal failure"),
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.StopParseDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error code for service error")
|
|
}
|
|
}
|
|
|
|
func TestStopParseDocumentsHandler_NotAccessible(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("POST", "/api/v1/datasets/ds-1/documents/stop", `{"document_ids": ["doc-1"]}`)
|
|
// Use a user that doesn't have access to ds-1
|
|
c.Set("user_id", "other-user")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.StopParseDocuments(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
code, _ := resp["code"].(float64)
|
|
if code == float64(common.CodeSuccess) {
|
|
t.Fatal("expected error for no authorization")
|
|
}
|
|
}
|
|
|
|
func TestStopIngestionTasksHandler_InvalidTransitionReturnsConflict(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
stopIngestionTaskErr: &service.InvalidTaskTransitionError{TaskID: "task-1", From: common.CREATED, To: common.COMPLETED},
|
|
}
|
|
h := &DocumentHandler{documentService: fake}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`)
|
|
h.StopIngestionTasks(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeConflict) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeConflict, resp["code"])
|
|
}
|
|
if !strings.Contains(resp["message"].(string), "task-1") {
|
|
t.Fatalf("expected task id in message, got %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestStopIngestionTasksHandler_TaskNotFoundReturnsNotFound(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
stopIngestionTaskErr: common.ErrTaskNotFound,
|
|
}
|
|
h := &DocumentHandler{documentService: fake}
|
|
|
|
c, w := setupGinContextWithUser("PUT", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`)
|
|
h.StopIngestionTasks(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeNotFound) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"])
|
|
}
|
|
}
|
|
|
|
func TestRemoveIngestionTasksHandler_TaskNotFoundReturnsNotFound(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
removeIngestionTaskErr: common.ErrTaskNotFound,
|
|
}
|
|
h := &DocumentHandler{documentService: fake}
|
|
|
|
c, w := setupGinContextWithUser("DELETE", "/api/v1/datasets/ds-1/ingestion/tasks", `{"tasks":["task-1"]}`)
|
|
h.RemoveIngestionTasks(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if resp["code"] != float64(common.CodeNotFound) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"])
|
|
}
|
|
}
|
|
|
|
func TestMetadataSummaryByDataset_Success(t *testing.T) {
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
|
|
gin.SetMode(gin.TestMode)
|
|
|
|
fake := &fakeDocumentService{
|
|
metadataSummary: map[string]interface{}{
|
|
"author": map[string]interface{}{
|
|
"type": "string",
|
|
"values": []interface{}{
|
|
[]interface{}{"alice", 2},
|
|
},
|
|
},
|
|
},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/datasets/ds-1/metadata/summary?doc_ids=doc-1,doc-2", "")
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: "ds-1"}}
|
|
|
|
h.MetadataSummaryByDataset(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if fake.metadataKBID != "ds-1" {
|
|
t.Fatalf("expected kbID ds-1, got %q", fake.metadataKBID)
|
|
}
|
|
if len(fake.metadataDocIDs) != 2 || fake.metadataDocIDs[0] != "doc-1" || fake.metadataDocIDs[1] != "doc-2" {
|
|
t.Fatalf("unexpected docIDs: %#v", fake.metadataDocIDs)
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("failed to unmarshal response: %v", err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected code 0, got %v: %v", resp["code"], resp)
|
|
}
|
|
data := resp["data"].(map[string]interface{})
|
|
summary := data["summary"].(map[string]interface{})
|
|
author := summary["author"].(map[string]interface{})
|
|
if author["type"] != "string" {
|
|
t.Fatalf("expected author type string, got %v", author["type"])
|
|
}
|
|
}
|
|
|
|
func TestGetThumbnail_Success(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
fake := &fakeDocumentService{
|
|
thumbnails: map[string]string{
|
|
"doc-1": "/api/v1/documents/doc-1/thumbnail",
|
|
"doc-2": "",
|
|
},
|
|
}
|
|
h := &DocumentHandler{
|
|
documentService: fake,
|
|
}
|
|
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/thumbnails?doc_ids=doc-1&doc_ids=doc-2", "")
|
|
|
|
h.GetThumbnail(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if len(fake.thumbnailDocIDs) != 2 || fake.thumbnailDocIDs[0] != "doc-1" || fake.thumbnailDocIDs[1] != "doc-2" {
|
|
t.Fatalf("unexpected docIDs: %#v", fake.thumbnailDocIDs)
|
|
}
|
|
if fake.thumbnailUserID != "user-1" {
|
|
t.Fatalf("unexpected userID: %s", fake.thumbnailUserID)
|
|
}
|
|
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("failed to unmarshal response: %v", err)
|
|
}
|
|
if resp["code"] != float64(common.CodeSuccess) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeSuccess, resp["code"])
|
|
}
|
|
data := resp["data"].(map[string]interface{})
|
|
if data["doc-1"] != "/api/v1/documents/doc-1/thumbnail" {
|
|
t.Fatalf("unexpected thumbnail for doc-1: %v", data["doc-1"])
|
|
}
|
|
if data["doc-2"] != "" {
|
|
t.Fatalf("unexpected thumbnail for doc-2: %v", data["doc-2"])
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentImagePassesPrincipalAndDisablesCaching(t *testing.T) {
|
|
fake := &fakeDocumentService{imageData: []byte("\x89PNG\r\n\x1a\nimage")}
|
|
h := &DocumentHandler{documentService: fake}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/doc-1/images/imagetemps-page-1.png", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}, {Key: "image_id", Value: "imagetemps-page-1.png"}}
|
|
|
|
h.GetDocumentImageForDocument(c)
|
|
|
|
if w.Code != http.StatusOK || w.Header().Get("Content-Type") != "image/png" || w.Header().Get("Cache-Control") != "no-store" {
|
|
t.Fatalf("unexpected response: status=%d content-type=%q cache=%q", w.Code, w.Header().Get("Content-Type"), w.Header().Get("Cache-Control"))
|
|
}
|
|
if fake.imageUserID != "user-1" || fake.imageDocumentID != "doc-1" || fake.imageID != "imagetemps-page-1.png" {
|
|
t.Fatalf("service arguments: user=%q doc=%q image=%q", fake.imageUserID, fake.imageDocumentID, fake.imageID)
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentImageRequiresPrincipal(t *testing.T) {
|
|
fake := &fakeDocumentService{imageData: []byte("\x89PNG\r\n\x1a\nimage")}
|
|
h := &DocumentHandler{documentService: fake}
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/documents/doc-1/images/imagetemps-page-1.png", nil)
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}, {Key: "image_id", Value: "imagetemps-page-1.png"}}
|
|
|
|
h.GetDocumentImageForDocument(c)
|
|
|
|
var response map[string]any
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if response["code"] == float64(common.CodeSuccess) || fake.imageID != "" {
|
|
t.Fatalf("anonymous request reached image service: response=%s image=%q", w.Body.String(), fake.imageID)
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentImageRejectsNonImagePayload(t *testing.T) {
|
|
fake := &fakeDocumentService{imageData: []byte("%PDF-1.7")}
|
|
h := &DocumentHandler{documentService: fake}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/images/anything.png", "")
|
|
c.Params = gin.Params{{Key: "image_id", Value: "anything.png"}}
|
|
|
|
h.GetDocumentImage(c)
|
|
|
|
var response map[string]any
|
|
if err := json.Unmarshal(w.Body.Bytes(), &response); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if response["code"] != float64(common.CodeDataError) {
|
|
t.Fatalf("unexpected response: %s", w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestDocumentImageContentTypeSupportedFormats(t *testing.T) {
|
|
tests := map[string]struct {
|
|
data []byte
|
|
want string
|
|
}{
|
|
"png": {[]byte("\x89PNG\r\n\x1a\n"), "image/png"},
|
|
"jpeg": {[]byte{0xff, 0xd8, 0xff}, "image/jpeg"},
|
|
"gif": {[]byte("GIF89a"), "image/gif"},
|
|
"webp": {[]byte("RIFF\x00\x00\x00\x00WEBP"), "image/webp"},
|
|
"bmp": {[]byte("BM"), "image/bmp"},
|
|
}
|
|
for name, test := range tests {
|
|
t.Run(name, func(t *testing.T) {
|
|
if got := documentImageContentType(test.data); got != test.want {
|
|
t.Fatalf("documentImageContentType() = %q, want %q", got, test.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentArtifact_Success(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/test.txt", "")
|
|
c.Params = gin.Params{{Key: "filename", Value: "test.txt"}}
|
|
|
|
h.GetDocumentArtifact(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
if w.Header().Get("Content-Type") != "text/plain" {
|
|
t.Fatalf("unexpected content type: %s", w.Header().Get("Content-Type"))
|
|
}
|
|
if w.Body.String() != "artifact content" {
|
|
t.Fatalf("unexpected body: %s", w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentArtifact_NotFound(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/error.txt", "")
|
|
c.Params = gin.Params{{Key: "filename", Value: "error.txt"}}
|
|
|
|
h.GetDocumentArtifact(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeDataError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeDataError, resp["code"])
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentArtifact_UnexpectedError(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/artifact/unexpected.txt", "")
|
|
c.Params = gin.Params{{Key: "filename", Value: "unexpected.txt"}}
|
|
|
|
h.GetDocumentArtifact(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeExceptionError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeExceptionError, resp["code"])
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentPreview_Success(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/doc-1/preview", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "doc-1"}}
|
|
|
|
h.GetDocumentPreview(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
if w.Header().Get("Content-Type") != "text/plain" {
|
|
t.Fatalf("unexpected content type: %s", w.Header().Get("Content-Type"))
|
|
}
|
|
if w.Body.String() != "preview content" {
|
|
t.Fatalf("unexpected body: %s", w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentPreview_NotFound(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/not-found/preview", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "not-found"}}
|
|
|
|
h.GetDocumentPreview(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeDataError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeDataError, resp["code"])
|
|
}
|
|
if resp["message"] != "document not found" {
|
|
t.Fatalf("expected message %q, got %v", "document not found", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentPreview_EmptyFile(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/empty-file/preview", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "empty-file"}}
|
|
|
|
h.GetDocumentPreview(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeDataError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeDataError, resp["code"])
|
|
}
|
|
if resp["message"] != "This file is empty." {
|
|
t.Fatalf("expected message %q, got %v", "This file is empty.", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestGetDocumentPreview_StorageErrorGenericMessage(t *testing.T) {
|
|
gin.SetMode(gin.TestMode)
|
|
h := &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
}
|
|
c, w := setupGinContextWithUser("GET", "/api/v1/documents/storage-error/preview", "")
|
|
c.Params = gin.Params{{Key: "id", Value: "storage-error"}}
|
|
|
|
h.GetDocumentPreview(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeServerError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeServerError, resp["code"])
|
|
}
|
|
msg, _ := resp["message"].(string)
|
|
if msg != "Failed to load document preview" {
|
|
t.Fatalf("expected generic message, got %v", resp["message"])
|
|
}
|
|
// The storage detail (bucket/key, transport error) must stay in the
|
|
// server log, not in the client-visible message.
|
|
if strings.Contains(msg, "connection refused") || strings.Contains(msg, "b/k") {
|
|
t.Fatalf("storage detail leaked to client: %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
// downloadHandlerWithAccessDB wires the real dataset permission check against
|
|
// the seeded in-memory DB (ds-1 is owned by tenant-1 / user-1) so the tests
|
|
// exercise the tenant gate on the download route.
|
|
func downloadHandlerWithAccessDB(t *testing.T) *DocumentHandler {
|
|
t.Helper()
|
|
db := setupHandlerAccessDB(t)
|
|
orig := dao.DB
|
|
dao.DB = db
|
|
t.Cleanup(func() { dao.DB = orig })
|
|
return &DocumentHandler{
|
|
documentService: &fakeDocumentService{},
|
|
datasetService: dataset.NewDatasetService(),
|
|
}
|
|
}
|
|
|
|
func downloadContextAs(userID, datasetID, docID string) (*gin.Context, *httptest.ResponseRecorder) {
|
|
gin.SetMode(gin.TestMode)
|
|
w := httptest.NewRecorder()
|
|
c, _ := gin.CreateTestContext(w)
|
|
c.Request = httptest.NewRequest("GET", "/api/v1/datasets/"+datasetID+"/documents/"+docID, nil)
|
|
c.Set("user", &entity.User{ID: userID})
|
|
c.Set("user_id", userID)
|
|
c.Params = gin.Params{{Key: "dataset_id", Value: datasetID}, {Key: "document_id", Value: docID}}
|
|
return c, w
|
|
}
|
|
|
|
func TestDownloadDocument_Success(t *testing.T) {
|
|
h := downloadHandlerWithAccessDB(t)
|
|
c, w := downloadContextAs("user-1", "ds-1", "doc-1")
|
|
|
|
h.DownloadDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
if w.Header().Get("Content-Type") != "application/pdf" {
|
|
t.Fatalf("unexpected content type: %s", w.Header().Get("Content-Type"))
|
|
}
|
|
if w.Body.String() != "document data" {
|
|
t.Fatalf("unexpected body: %s", w.Body.String())
|
|
}
|
|
}
|
|
|
|
func TestDownloadDocument_NotFound(t *testing.T) {
|
|
h := downloadHandlerWithAccessDB(t)
|
|
c, w := downloadContextAs("user-1", "ds-1", "not-found")
|
|
|
|
h.DownloadDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeNotFound) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"])
|
|
}
|
|
if resp["message"] != "Resource not found" {
|
|
t.Fatalf("unexpected hidden-resource message: %v", resp["message"])
|
|
}
|
|
}
|
|
|
|
func TestDownloadDocument_ServiceErrorReturnsServerError(t *testing.T) {
|
|
h := downloadHandlerWithAccessDB(t)
|
|
h.documentService = &fakeDocumentService{downloadErr: errors.New("database connection failed")}
|
|
c, w := downloadContextAs("user-1", "ds-1", "doc-1")
|
|
|
|
h.DownloadDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200 response envelope, got %d", w.Code)
|
|
}
|
|
var resp map[string]interface{}
|
|
if err := json.Unmarshal(w.Body.Bytes(), &resp); err != nil {
|
|
t.Fatalf("decode response: %v", err)
|
|
}
|
|
if resp["code"] != float64(common.CodeServerError) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeServerError, resp["code"])
|
|
}
|
|
if resp["message"] != "Failed to download document" {
|
|
t.Fatalf("unexpected client-visible error: %v", resp["message"])
|
|
}
|
|
if strings.Contains(w.Body.String(), "database connection failed") {
|
|
t.Fatal("database error detail leaked to the client")
|
|
}
|
|
}
|
|
|
|
// A logged-in user who is not a member of the dataset's tenant must not be
|
|
// able to download its documents, and must not learn whether they exist.
|
|
func TestDownloadDocument_ForeignUserRejected(t *testing.T) {
|
|
h := downloadHandlerWithAccessDB(t)
|
|
c, w := downloadContextAs("user-2", "ds-1", "doc-1")
|
|
|
|
h.DownloadDocument(c)
|
|
|
|
if w.Code != http.StatusOK {
|
|
t.Fatalf("expected 200 envelope, got %d: %s", w.Code, w.Body.String())
|
|
}
|
|
if strings.Contains(w.Body.String(), "document data") {
|
|
t.Fatalf("document bytes must not be served to a foreign user")
|
|
}
|
|
var resp map[string]interface{}
|
|
json.Unmarshal(w.Body.Bytes(), &resp)
|
|
if resp["code"] != float64(common.CodeNotFound) {
|
|
t.Fatalf("expected code %d, got %v", common.CodeNotFound, resp["code"])
|
|
}
|
|
if resp["message"] != "Resource not found" {
|
|
t.Fatalf("unexpected hidden-resource message: %v", resp["message"])
|
|
}
|
|
}
|