Files
ragflow/internal/service/dataset/update_test.go
T

1626 lines
58 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 dataset
import (
"encoding/json"
"reflect"
"strings"
"testing"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/entity"
"ragflow/internal/service"
"gorm.io/gorm"
)
func TestDatasetServiceUpdateDatasetUpdatesFields(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateCanvas(t, "abcdef0123456789abcdef0123456789", "tenant-1")
name := " Renamed Dataset "
description := "updated description"
language := "English"
permission := string(entity.TenantPermissionTeam)
chunkMethod := string(entity.ParserTypeBook)
embeddingModel := "BAAI/bge-large-zh-v1.5@Builtin"
parseType := 1
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &name,
Description: &description,
Language: &language,
Permission: &permission,
ParserID: &chunkMethod,
ParseType: &parseType,
EmbeddingModel: &embeddingModel,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["name"] != "Renamed Dataset" {
t.Fatalf("expected trimmed name in response, got %#v", result["name"])
}
if result["parser_id"] != chunkMethod {
t.Fatalf("expected parser_id %q, got %#v", chunkMethod, result["parser_id"])
}
if result["embedding_model"] != embeddingModel {
t.Fatalf("expected embedding model %q, got %#v", embeddingModel, result["embedding_model"])
}
if connectors, ok := result["connectors"].([]*dao.ConnectorDatasetListItem); !ok || len(connectors) != 0 {
t.Fatalf("expected empty connector list, got %#v", result["connectors"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.Name != "Renamed Dataset" {
t.Fatalf("expected persisted trimmed name, got %q", persisted.Name)
}
if persisted.Description == nil || *persisted.Description != description {
t.Fatalf("expected persisted description, got %#v", persisted.Description)
}
if persisted.Permission != permission {
t.Fatalf("expected persisted permission %q, got %q", permission, persisted.Permission)
}
if persisted.ParserID != chunkMethod {
t.Fatalf("expected parser id %q, got %q", chunkMethod, persisted.ParserID)
}
if persisted.EmbdID != embeddingModel {
t.Fatalf("expected embd id %q, got %q", embeddingModel, persisted.EmbdID)
}
// parser_config stores DSL runtime component params directly.
if _, ok := persisted.ParserConfig["Parser:HipSignsRhyme"].(map[string]interface{}); !ok {
t.Fatalf("expected Parser:HipSignsRhyme in parser_config, got %#v", persisted.ParserConfig)
}
}
func TestUpdateDataset_ClearsAvatarWithEmptyString(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
svc := testDatasetUpdateService(t)
initial := "data:image/png;base64,iVBORw0KGgo="
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{Avatar: &initial}); err != nil || code != common.CodeSuccess {
t.Fatalf("set avatar failed: code=%d err=%v", code, err)
}
empty := ""
result, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{Avatar: &empty})
if err != nil {
t.Fatalf("clear avatar failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["avatar"] != "" {
t.Fatalf("expected empty avatar in response, got %#v", result["avatar"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.Avatar == nil || *persisted.Avatar != "" {
t.Fatalf("expected persisted empty avatar, got %#v", persisted.Avatar)
}
}
// TestUpdateDataset_ParentChildConfigReachesGeneralChunker verifies an edit to
// the dataset setting updates the runtime chunker parameter, not only the UI
// payload retained in parser_config.
func TestUpdateDataset_ParentChildConfigReachesGeneralChunker(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"GeneralChunker:SixApplesFall": map[string]interface{}{
"parent_child": map[string]interface{}{
"use_parent_child": true,
"children_delimiter": "|",
},
},
},
})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDataset err=%v code=%d", err, code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), db, "kb-1")
if err != nil {
t.Fatalf("get updated dataset: %v", err)
}
chunker, ok := persisted.ParserConfig["GeneralChunker:SixApplesFall"].(map[string]interface{})
if !ok {
t.Fatalf("general chunker params = %#v", persisted.ParserConfig["GeneralChunker:SixApplesFall"])
}
if got, ok := chunker["children_delimiters"].([]interface{}); !ok || len(got) != 1 || got[0] != "|" {
t.Fatalf("children_delimiters = %#v, want [\"|\"]", chunker["children_delimiters"])
}
parentChild, ok := chunker["parent_child"].(map[string]interface{})
if !ok || parentChild["use_parent_child"] != true || parentChild["children_delimiter"] != "|" {
t.Fatalf("chunker parent_child = %#v, want persisted public setting", chunker["parent_child"])
}
if _, ok := persisted.ParserConfig["parent_child"]; ok {
t.Fatalf("top-level flat parent_child should be absent, got %#v", persisted.ParserConfig["parent_child"])
}
}
func TestUpdateDatasetPreservesParentChildChunkerRuntimeConfig(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
existingConfig := entity.JSONMap{
"GeneralChunker:SixApplesFall": map[string]interface{}{
"parent_child": map[string]interface{}{
"use_parent_child": true,
"children_delimiter": "|",
},
},
}
if err := db.Model(&entity.Knowledgebase{}).Where("id = ?", "kb-1").Update("parser_config", existingConfig).Error; err != nil {
t.Fatalf("seed parent-child parser config: %v", err)
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
// Component-scoped override (flat parser-level keys are rejected now).
ParserConfig: map[string]interface{}{
"GeneralChunker:SixApplesFall": map[string]interface{}{"chunk_token_size": float64(256)},
},
})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDataset err=%v code=%d", err, code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), db, "kb-1")
if err != nil {
t.Fatalf("load updated dataset: %v", err)
}
chunker, ok := persisted.ParserConfig["GeneralChunker:SixApplesFall"].(map[string]interface{})
if !ok {
t.Fatalf("general chunker params = %#v", persisted.ParserConfig["GeneralChunker:SixApplesFall"])
}
if got, ok := chunker["children_delimiters"].([]interface{}); !ok || len(got) != 1 || got[0] != "|" {
t.Fatalf("children_delimiters = %#v, want [\"|\"]", chunker["children_delimiters"])
}
}
func TestUpdateDataset_RejectsSimultaneousParserIDAndPipelineID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
chunkMethod := "book"
pipelineID := "abcdef0123456789abcdef0123456789"
ctx := t.Context()
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
PipelineID: &pipelineID,
})
if err == nil {
t.Fatal("expected error when both parser_id and pipeline_id are set without parse_type")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != "parse_type is required" {
t.Fatalf("expected 'parse_type is required', got: %v", err)
}
}
func TestUpdateDataset_ParseTypeBuiltinClearsPipelineID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
seedDatasetUpdateCanvas(t, "abcdef0123456789abcdef0123456789", "tenant-1",
datasetUpdateCanvasDSL("Parser:HipSignsRhyme", "chunk_token_num"))
chunkMethod := "book"
parseType := 1
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
// parse_type=1 selects builtin mode → only parser_id should be set.
if result["parser_id"] != chunkMethod {
t.Fatalf("expected parser_id %q, got %#v", chunkMethod, result["parser_id"])
}
if pid, ok := result["pipeline_id"]; ok && pid != nil && pid != "" {
t.Fatalf("expected pipeline_id to be cleared for BuiltIn mode, got %#v", pid)
}
}
// TestUpdateDataset_ParseTypePipelineRejectsParserID locks in the new explicit
// contract: parse_type=2 (Pipeline) must not carry a parser_id. The previous
// lenient behavior silently ignored the contradictory id; now it is rejected so
// a malformed request can never pick the wrong mode.
func TestUpdateDataset_ParseTypePipelineRejectsParserID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
chunkMethod := "book"
pipelineID := "ABCDEF0123456789ABCDEF0123456789"
parseType := 2
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
PipelineID: &pipelineID,
ParseType: &parseType,
})
if err == nil {
t.Fatal("expected error rejecting parser_id in Pipeline mode")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != "parser_id must not be set when parse_type is Pipeline" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestUpdateDataset_PipelineSwitchKeepsSubmittedParserConfig(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
pipelineID := "abcdef0123456789abcdef0123456789"
parseType := 2
dslJSON, err := json.Marshal(map[string]any{
"components": map[string]any{
"Parser:CustomP": map[string]any{
"obj": map[string]any{
"component_name": "Parser",
"params": map[string]any{
"pdf": map[string]any{"remove_toc": false},
},
},
},
},
})
if err != nil {
t.Fatalf("marshal canvas DSL: %v", err)
}
seedDatasetUpdateCanvas(t, pipelineID, "tenant-1", dslJSON)
// C5/G3: a Pipeline-mode switch now requires an explicit version (the empty
// version is no longer auto-backfilled). Seed one and pass it.
verID := "version0000000000000000000000cc"
seedCanvasVersion(t, db, pipelineID, verID, 0)
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
PipelineID: &pipelineID,
ParseType: &parseType,
PipelineVersionID: &verID,
ParserConfig: map[string]interface{}{
"Parser:CustomP": map[string]interface{}{
"pdf": map[string]interface{}{"remove_toc": true},
},
},
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
// A pipeline (canvas) switch with a submitted parser_config keeps the
// submitted values, reconciled against the new canvas DSL by
// BuildParserConfig: the settings form resets itself to the new defaults
// on a switch, so the submitted config carries the user's edits of the new
// pipeline and must not be discarded. remove_toc=true must survive on top
// of the new canvas DSL default (false).
parserParams, ok := persisted.ParserConfig["Parser:CustomP"].(map[string]interface{})
if !ok {
t.Fatalf("expected Parser:CustomP in parser_config, got %#v", persisted.ParserConfig)
}
pdfParams, ok := parserParams["pdf"].(map[string]interface{})
if !ok {
t.Fatalf("expected pdf parser config, got %#v", parserParams)
}
if removeTOC, ok := pdfParams["remove_toc"].(bool); !ok || !removeTOC {
t.Fatalf("remove_toc = %#v, want true (submitted value kept on pipeline switch)", pdfParams["remove_toc"])
}
}
// TestUpdateDataset_ParseTypePipelineCleansConfigAgainstCanvas verifies that
// with parse_type=2 and a dirty parser_id, parser_config is still cleaned
// against the canvas DSL (not the builtin DSL). This locks in the
// ResolveParseMode contract on the dataset path: mode is authoritative and
// ignores the conflicting parser_id.
func TestUpdateDataset_ParseTypePipelineCleansConfigAgainstCanvas(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original",
"abcdef0123456789abcdef0123456789")
seedDatasetUpdateCanvas(t, "abcdef0123456789abcdef0123456789", "tenant-1",
datasetUpdateCanvasDSL("Parser:CustomP", "chunk_token_num"))
// parse_type=2 must clean the config against the canvas DSL. No parser_id is
// sent in pipeline mode (the new contract forbids a contradictory id).
pipelineID := "abcdef0123456789abcdef0123456789"
parseType := 2
override := map[string]interface{}{
"Parser:CustomP": map[string]interface{}{"chunk_token_num": float64(256)},
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
PipelineID: &pipelineID,
ParseType: &parseType,
ParserConfig: override,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
// Read back the persisted KB and verify parser_config was cleaned against
// the canvas DSL: it must carry the canvas component (Parser:CustomP), not
// a builtin-book component structure.
var kb entity.Knowledgebase
if err := dao.DB.First(&kb, "id = ?", "kb-1").Error; err != nil {
t.Fatalf("read back kb: %v", err)
}
cleaned, ok := kb.ParserConfig["Parser:CustomP"]
if !ok {
t.Fatalf("parser_config = %v, want Parser:CustomP component from canvas DSL", kb.ParserConfig)
}
cpn, ok := cleaned.(map[string]any)
if !ok {
t.Fatalf("Parser:CustomP = %T, want map", cleaned)
}
if v, _ := cpn["chunk_token_num"].(float64); v != 256 {
t.Fatalf("chunk_token_num = %v, want 256 (override applied on canvas defaults)", cpn["chunk_token_num"])
}
}
// TestUpdateDatasetRejectsInvalidPages verifies the fail-fast contract for the
// "pages" range: an invalid range (from>to) aborts the request with
// CodeDataError instead of being silently dropped or persisted.
func TestUpdateDatasetRejectsInvalidPages(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
chunkMethod := "manual"
parseType := 1
ctx := t.Context()
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{
"pages": []interface{}{[]interface{}{float64(5), float64(3)}}, // from>to invalid
},
},
},
})
if err == nil {
t.Fatalf("expected error for invalid pages range, got nil")
}
if code != common.CodeDataError {
t.Fatalf("code = %v, want CodeDataError", code)
}
}
func TestDatasetServiceGetDatasetReturnsEmptyConnectorList(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
datasetID := "11111111111141118111111111111111"
insertDatasetUpdateKB(t, datasetID, "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).GetDataset(ctx, "11111111-1111-4111-8111-111111111111")
if err != nil {
t.Fatalf("GetDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
connectors, ok := result["connectors"].([]*dao.ConnectorDatasetListItem)
if !ok {
t.Fatalf("expected connector list, got %#v", result["connectors"])
}
if connectors == nil {
t.Fatal("expected empty connector list, got nil")
}
if len(connectors) != 0 {
t.Fatalf("expected empty connector list, got %#v", connectors)
}
}
func TestDatasetServiceUpdateDatasetRejectsMissingDataset(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
name := "Renamed"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "missing-kb", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected missing dataset error")
}
if code != common.CodeNotFound {
t.Fatalf("expected not found code, got %d", code)
}
if err.Error() != "Resource not found" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetRejectsNonOwner(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ensureDatasetTestMembership(t, "tenant-2", "tenant-2", "owner")
name := "Renamed"
ctx := withDatasetTestContext(t.Context(), "tenant-2", "tenant-2")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected permission error")
}
if code != common.CodeNotFound {
t.Fatalf("expected not found code, got %d", code)
}
if err.Error() != "Resource not found" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetRejectsTeamMemberPermissionChange(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "owner-1", "Original")
if err := dao.DB.Model(&entity.Knowledgebase{}).
Where("id = ?", "kb-1").
Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil {
t.Fatalf("update kb permission: %v", err)
}
insertDatasetUpdateTeamMember(t, "user-1", "owner-1")
permission := string(entity.TenantPermissionMe)
ctx := withDatasetTestContext(t.Context(), "user-1", "owner-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Permission: &permission,
})
if err == nil {
t.Fatal("expected shared tenant member to be denied changing dataset permission")
}
if code == common.CodeSuccess {
t.Fatalf("expected a permission error code, got success (%d)", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get dataset: %v", err)
}
if persisted.Permission != string(entity.TenantPermissionTeam) {
t.Fatalf("expected permission to remain %q, got %q", entity.TenantPermissionTeam, persisted.Permission)
}
}
func TestDatasetServiceUpdateDatasetAllowsTeamMemberToResubmitUnchangedPermission(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "owner-1", "Original")
if err := dao.DB.Model(&entity.Knowledgebase{}).
Where("id = ?", "kb-1").
Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil {
t.Fatalf("update kb permission: %v", err)
}
insertDatasetUpdateTeamMember(t, "user-1", "owner-1")
grantDatasetTestResource(t, "owner-1", "kb", "kb-1", "user-1", "write")
name := "Renamed by team member"
permission := string(entity.TenantPermissionTeam)
result, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "owner-1"), "kb-1", service.UpdateDatasetRequest{
Name: &name,
Permission: &permission,
})
if err != nil {
t.Fatalf("expected unchanged permission to be accepted, got: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["name"] != name {
t.Fatalf("expected updated name %q, got %#v", name, result["name"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "user-1", "owner-1"), db, "kb-1")
if err != nil {
t.Fatalf("get dataset: %v", err)
}
if persisted.Name != name {
t.Fatalf("expected persisted name %q, got %q", name, persisted.Name)
}
if persisted.Permission != permission {
t.Fatalf("expected permission to remain %q, got %q", permission, persisted.Permission)
}
}
func TestDatasetServiceUpdateDatasetResolvesEmbeddingInDatasetTenant(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
if err := dao.DB.Model(&entity.Knowledgebase{}).
Where("id = ?", "kb-1").
Update("permission", string(entity.TenantPermissionTeam)).Error; err != nil {
t.Fatalf("update kb permission: %v", err)
}
insertDatasetUpdateTeamMember(t, "user-1", "tenant-1")
grantDatasetTestResource(t, "tenant-1", "kb", "kb-1", "user-1", "write")
insertDatasetUpdateModelProvider(t, "provider-1", "tenant-1", "ZHIPU-AI")
insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test")
embeddingModelID := "aabbccdd11223344aabbccdd11223344"
insertDatasetUpdateTenantModel(t, embeddingModelID, "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding))
enableDatasetTestModelAccess(t, "tenant-1", "user-1", "instance-1")
permission := string(entity.TenantPermissionTeam)
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "user-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
Permission: &permission,
EmbeddingModel: &embeddingModelID,
})
if err != nil {
t.Fatalf("expected shared dataset embedding model to resolve in its owning tenant, got: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "user-1", "tenant-1"), db, "kb-1")
if err != nil {
t.Fatalf("get dataset: %v", err)
}
if persisted.TenantEmbdID == nil || *persisted.TenantEmbdID != embeddingModelID {
t.Fatalf("expected tenant_embd_id %q, got %#v", embeddingModelID, persisted.TenantEmbdID)
}
}
func TestDatasetServiceUpdateDatasetValidatesName(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
name := " "
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected name validation error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != "String should have at least 1 character" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetRejectsDuplicateName(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Existing")
name := "Existing"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{Name: &name})
if err == nil {
t.Fatal("expected duplicate name error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if !strings.Contains(err.Error(), "already exists") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetRejectsNoPropertiesModified(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{})
if err == nil {
t.Fatal("expected no-op update error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != "no properties were modified" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetLinksConnectors(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateConnector(t, "connector-1", "tenant-1")
autoParse := "0"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}},
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
connectors, ok := result["connectors"].([]*dao.ConnectorDatasetListItem)
if !ok {
t.Fatalf("expected connector list, got %#v", result["connectors"])
}
if len(connectors) != 1 {
t.Fatalf("expected one connector, got %d", len(connectors))
}
if connectors[0].ID != "connector-1" || connectors[0].AutoParse != autoParse {
t.Fatalf("unexpected connector: %#v", connectors[0])
}
var link entity.Connector2Kb
if err := dao.DB.Where("kb_id = ? AND connector_id = ?", "kb-1", "connector-1").First(&link).Error; err != nil {
t.Fatalf("expected connector link: %v", err)
}
if link.AutoParse != autoParse {
t.Fatalf("expected auto_parse %q, got %q", autoParse, link.AutoParse)
}
}
func TestDatasetServiceUpdateDatasetRejectsCrossTenantConnector(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateConnector(t, "connector-1", "tenant-2")
autoParse := "0"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &[]service.DatasetConnectorRequest{{ID: "connector-1", AutoParse: autoParse}},
})
if err == nil {
t.Fatal("expected connector tenant mismatch error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if !dao.IsConnectorNotAccessibleErr(err) {
t.Fatalf("unexpected error: %v", err)
}
var count int64
if err := dao.DB.Model(&entity.Connector2Kb{}).
Where("kb_id = ? AND connector_id = ?", "kb-1", "connector-1").
Count(&count).Error; err != nil {
t.Fatalf("count connector link: %v", err)
}
if count != 0 {
t.Fatalf("expected no connector link, got %d", count)
}
}
func TestDatasetServiceUpdateDatasetAcceptsProviderInstanceEmbedding(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateModelProvider(t, "provider-1", "tenant-1", "ZHIPU-AI")
insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test")
insertDatasetUpdateTenantModel(t, "model-1", "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding))
embeddingModel := "embedding-2@test@ZHIPU-AI"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embeddingModel,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["embedding_model"] != embeddingModel {
t.Fatalf("expected embedding model %q, got %#v", embeddingModel, result["embedding_model"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.EmbdID != embeddingModel {
t.Fatalf("expected persisted embedding model %q, got %q", embeddingModel, persisted.EmbdID)
}
}
func TestDatasetServiceUpdateDatasetAcceptsEmbeddingModelID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateModelProvider(t, "provider-1", "tenant-1", "ZHIPU-AI")
insertDatasetUpdateModelInstance(t, "instance-1", "provider-1", "test")
insertDatasetUpdateTenantModel(t, "aabbccdd11223344aabbccdd11223344", "provider-1", "instance-1", "embedding-2", int(entity.ModelTypeEmbedding))
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
embeddingModelID := "aabbccdd11223344aabbccdd11223344"
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embeddingModelID,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["embedding_model"] != embeddingModelID {
t.Fatalf("expected embedding model %q, got %#v", embeddingModelID, result["embedding_model"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.EmbdID != embeddingModelID {
t.Fatalf("expected persisted embedding model %q, got %q", embeddingModelID, persisted.EmbdID)
}
if persisted.TenantEmbdID == nil || *persisted.TenantEmbdID != embeddingModelID {
t.Fatalf("expected persisted tenant_embd_id %q, got %#v", embeddingModelID, persisted.TenantEmbdID)
}
}
func TestDatasetServiceUpdateDatasetRejectsEmptyConnectorID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
connectors := []service.DatasetConnectorRequest{{ID: " "}}
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Connectors: &connectors,
})
if err == nil {
t.Fatal("expected connector validation error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != "connector id is required" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetRejectsInvalidEmbeddingModelFormat(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
cases := []struct {
name string
embeddingModel string
expectedMessage string
}{
{"empty", "", "embedding model identifier must follow <model_name>@<provider> format"},
{"whitespace", " ", "embedding model identifier must follow <model_name>@<provider> format"},
{"missing_at", "BAAI/bge-small-en-v1.5Builtin", "embedding model identifier must follow <model_name>@<provider> format"},
{"empty_model_name", "@Builtin", "both model_name and provider must be non-empty strings"},
{"empty_provider", "BAAI/bge-small-en-v1.5@", "both model_name and provider must be non-empty strings"},
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
svc := testDatasetUpdateService(t)
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
embdModel := tc.embeddingModel
_, code, err := svc.UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
EmbeddingModel: &embdModel,
})
if err == nil {
t.Fatal("expected embedding model format error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if err.Error() != tc.expectedMessage {
t.Fatalf("unexpected error: got %q, want %q", err.Error(), tc.expectedMessage)
}
})
}
}
func TestDatasetServiceUpdateDatasetRejectsDuplicateNameCaseInsensitive(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
insertDatasetUpdateKB(t, "kb-2", "tenant-1", "Existing")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
uppercaseName := "EXISTING"
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &uppercaseName,
})
if err == nil {
t.Fatal("expected case-insensitive duplicate name error")
}
if code != common.CodeDataError {
t.Fatalf("expected data error code, got %d", code)
}
if !strings.Contains(err.Error(), "already exists") {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceUpdateDatasetPreservesUnmodifiedFields(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
description := "original description"
language := "English"
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
dao.DB.Model(&entity.Knowledgebase{}).Where("id = ?", "kb-1").Updates(map[string]interface{}{
"description": description,
"language": language,
})
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
newName := "Renamed Only"
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &newName,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["name"] != newName {
t.Fatalf("expected updated name %q, got %#v", newName, result["name"])
}
if result["description"] != description {
t.Fatalf("expected description preserved, got %#v", result["description"])
}
if result["language"] != language {
t.Fatalf("expected language preserved, got %#v", result["language"])
}
if result["embedding_model"] != "BAAI/bge-large-zh-v1.5@Builtin" {
t.Fatalf("expected embedding_model preserved, got %#v", result["embedding_model"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.Name != newName {
t.Fatalf("expected persisted name %q, got %q", newName, persisted.Name)
}
if persisted.Description == nil || *persisted.Description != description {
t.Fatalf("expected persisted description %q, got %#v", description, persisted.Description)
}
}
func TestDatasetServiceUpdateDatasetPreservesParserConfigOnEmptyUpdate(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
dao.DB.Model(&entity.Knowledgebase{}).Where("id = ?", "kb-1").Update("parser_config", entity.JSONMap{
"chunk_token_num": float64(512),
"delimiter": "\n",
})
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
name := "Updated Name"
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
Name: &name,
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
if persisted.ParserConfig["chunk_token_num"] != float64(512) {
t.Fatalf("expected chunk_token_num preserved, got %#v", persisted.ParserConfig["chunk_token_num"])
}
if persisted.ParserConfig["delimiter"] != "\n" {
t.Fatalf("expected delimiter preserved, got %#v", persisted.ParserConfig["delimiter"])
}
}
func TestDatasetServiceDeleteDatasetsRejectsUnauthorizedID(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "11111111111141118111111111111111", "tenant-1", "Test")
ensureDatasetTestMembership(t, "tenant-2", "tenant-1", "normal")
ctx := withDatasetTestContext(t.Context(), "tenant-2", "tenant-1")
svc := NewDatasetService()
normalizedID := "11111111111141118111111111111111"
_, code, err := svc.DeleteDatasets(ctx, []string{normalizedID}, false)
if err == nil {
t.Fatal("expected unauthorized error")
}
if code != common.CodeForbidden {
t.Fatalf("expected forbidden code, got %d", code)
}
if err.Error() != "Permission denied" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceDeleteDatasetsRejectsAllUnauthorized(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
svc := NewDatasetService()
_, code, err := svc.DeleteDatasets(ctx, []string{"d94a8dc02c9711f0930f7fbc369eab6d"}, false)
if err == nil {
t.Fatal("expected unauthorized error")
}
if code != common.CodeForbidden {
t.Fatalf("expected forbidden code, got %d", code)
}
if err.Error() != "Permission denied" {
t.Fatalf("unexpected error: %v", err)
}
}
func TestDatasetServiceDeleteAllDatasetsRequiresTenantOwner(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Shared dataset")
insertDatasetUpdateTeamMember(t, "member-1", "tenant-1")
_, code, err := NewDatasetService().DeleteDatasets(
withDatasetTestContext(t.Context(), "member-1", "tenant-1"), nil, true,
)
if err == nil {
t.Fatal("expected delete-all to be denied for a non-owner tenant member")
}
if code != common.CodeForbidden {
t.Fatalf("expected forbidden code, got %d", code)
}
if err.Error() != "Permission denied" {
t.Fatalf("unexpected error: %v", err)
}
}
func setupDatasetUpdateTestDB(t *testing.T) *gorm.DB {
t.Helper()
db := setupServiceTestDB(t)
migrateDatasetUpdateTestTables(t, db)
return db
}
func migrateDatasetUpdateTestTables(t *testing.T, db *gorm.DB) {
t.Helper()
if err := db.AutoMigrate(
&entity.Connector{},
&entity.Connector2Kb{},
&entity.SyncLogs{},
&entity.TenantModelProvider{},
&entity.TenantModelInstance{},
&entity.TenantModel{},
&entity.UserCanvas{},
&entity.UserCanvasVersion{},
); err != nil {
t.Fatalf("failed to migrate dataset update tables: %v", err)
}
}
func testDatasetUpdateService(t *testing.T) *DatasetService {
t.Helper()
return &DatasetService{
kbDAO: dao.NewKnowledgebaseDAO(),
documentDAO: dao.NewDocumentDAO(),
connectorDAO: dao.NewConnectorDAO(),
}
}
func insertDatasetUpdateKB(t *testing.T, id, tenantID, name string) {
t.Helper()
kb := &entity.Knowledgebase{
ID: id,
TenantID: tenantID,
Name: name,
EmbdID: "BAAI/bge-large-zh-v1.5@Builtin",
CreatedBy: tenantID,
Permission: string(entity.TenantPermissionMe),
ParserID: string(entity.ParserTypeNaive),
ParserConfig: entity.JSONMap{"chunk_token_num": float64(128)},
Status: sptr(string(entity.StatusValid)),
}
if err := dao.DB.Create(kb).Error; err != nil {
t.Fatalf("insert test kb: %v", err)
}
ensureDatasetTestMembership(t, tenantID, tenantID, "owner")
}
func insertDatasetUpdateCanvas(t *testing.T, id, userID string) {
t.Helper()
canvas := &entity.UserCanvas{
ID: id,
UserID: userID,
DSL: entity.JSONMap{},
}
if err := dao.DB.Create(canvas).Error; err != nil {
t.Fatalf("insert test agent: %v", err)
}
}
func insertDatasetUpdateTeamMember(t *testing.T, userID, tenantID string) {
t.Helper()
ensureDatasetTestMembership(t, userID, tenantID, "normal")
}
func insertDatasetUpdateConnector(t *testing.T, id, tenantID string) {
t.Helper()
connector := &entity.Connector{
ID: id,
TenantID: tenantID,
Name: "Test Connector",
Source: "google_drive",
InputType: "oauth",
Config: entity.ConnectorConfig{"sync_deleted_files": false},
Status: string(entity.TaskStatusDone),
}
if err := dao.DB.Create(connector).Error; err != nil {
t.Fatalf("insert test connector: %v", err)
}
}
func insertDatasetUpdateModelProvider(t *testing.T, id, tenantID, providerName string) {
t.Helper()
provider := &entity.TenantModelProvider{
ID: id,
UserID: tenantID,
TenantID: tenantID,
ProviderName: providerName,
}
if err := dao.DB.Create(provider).Error; err != nil {
t.Fatalf("insert test model provider: %v", err)
}
}
func insertDatasetUpdateModelInstance(t *testing.T, id, providerID, instanceName string) {
t.Helper()
instance := &entity.TenantModelInstance{
ID: id,
ProviderID: providerID,
InstanceName: instanceName,
APIKey: "test-api-key",
Status: "active",
Extra: "{}",
}
if err := dao.DB.Create(instance).Error; err != nil {
t.Fatalf("insert test model instance: %v", err)
}
}
func insertDatasetUpdateTenantModel(t *testing.T, id, providerID, instanceID, modelName string, modelType int) {
t.Helper()
model := &entity.TenantModel{
ID: id,
ProviderID: providerID,
InstanceID: instanceID,
ModelName: modelName,
ModelType: modelType,
Status: "active",
Extra: "{}",
}
if err := dao.DB.Create(model).Error; err != nil {
t.Fatalf("insert test tenant model: %v", err)
}
}
// seedDatasetUpdateCanvas inserts a canvas row with the given DSL.
func seedDatasetUpdateCanvas(t *testing.T, id, userID string, dslJSON []byte) {
t.Helper()
var dslMap map[string]any
if err := json.Unmarshal(dslJSON, &dslMap); err != nil {
t.Fatalf("unmarshal seed dsl: %v", err)
}
canvas := &entity.UserCanvas{
ID: id,
UserID: userID,
Tags: "",
Permission: "me",
CanvasCategory: "agent_canvas",
DSL: entity.JSONMap(dslMap),
}
if err := dao.DB.Create(canvas).Error; err != nil {
t.Fatalf("seed agent: %v", err)
}
}
// datasetUpdateCanvasDSL builds a minimal canvas DSL declaring a Parser
// component with the given cpnID and param keys.
func datasetUpdateCanvasDSL(cpnID string, paramKeys ...string) []byte {
params := map[string]any{"outputs": map[string]any{}}
for _, k := range paramKeys {
params[k] = map[string]any{}
}
dsl := map[string]any{
"components": map[string]any{
cpnID: map[string]any{
"obj": map[string]any{"component_name": "Parser", "params": params},
},
},
}
raw, _ := json.Marshal(dsl)
return raw
}
// --- Step 3: component_params validation wired into UpdateDataset ---
// insertDatasetUpdateCanvasKB seeds a KB bound to a custom canvas pipeline
// (PipelineID set, ParserID empty) for the canvas validation tests.
func insertDatasetUpdateCanvasKB(t *testing.T, id, tenantID, name, pipelineID string) {
t.Helper()
pid := pipelineID
kb := &entity.Knowledgebase{
ID: id,
TenantID: tenantID,
Name: name,
EmbdID: "BAAI/bge-large-zh-v1.5@Builtin",
CreatedBy: tenantID,
Permission: string(entity.TenantPermissionMe),
ParserID: "",
PipelineID: &pid,
ParserConfig: entity.JSONMap{"chunk_token_num": float64(128)},
Status: sptr(string(entity.StatusValid)),
}
if err := dao.DB.Create(kb).Error; err != nil {
t.Fatalf("insert test canvas kb: %v", err)
}
ensureDatasetTestMembership(t, tenantID, tenantID, "owner")
}
func TestUpdateDataset_StripsUnknownParam_Builtin(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"no_such_param": 1,
},
},
})
if err != nil {
t.Fatalf("UpdateDataset should succeed (unknown params are stripped): %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
cp := map[string]interface{}(persisted.ParserConfig)
rhyme, ok := cp["Parser:HipSignsRhyme"].(map[string]interface{})
if !ok {
t.Fatalf("expected Parser:HipSignsRhyme persisted, got %#v", cp)
}
if _, exists := rhyme["no_such_param"]; exists {
t.Fatal("expected no_such_param to be stripped")
}
if result["parser_id"] != string(entity.ParserTypeGeneral) {
t.Fatalf("expected canonical parser_id %q, got %#v", entity.ParserTypeGeneral, result["parser_id"])
}
}
func TestUpdateDataset_AcceptsValidComponentParams_Builtin(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
},
},
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["parser_id"] != string(entity.ParserTypeGeneral) {
t.Fatalf("expected canonical parser_id %q, got %#v", entity.ParserTypeGeneral, result["parser_id"])
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
cp := map[string]interface{}(persisted.ParserConfig)
if len(cp) == 0 {
t.Fatalf("expected component_params persisted, got %#v", persisted.ParserConfig)
}
rhyme, ok := cp["Parser:HipSignsRhyme"].(map[string]interface{})
if !ok {
t.Fatalf("expected Parser:HipSignsRhyme persisted, got %#v", cp)
}
pdf, ok := rhyme["pdf"].(map[string]interface{})
if !ok || pdf["parse_method"] != "deepdoc" {
t.Fatalf("expected pdf setup persisted, got %#v", rhyme)
}
}
func TestUpdateDatasetUsesTenantModelIDForExtractorDefault(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
if err := db.Model(&entity.UserTenant{}).
Where("user_id = ? AND tenant_id = ?", "tenant-1", "tenant-1").
Updates(map[string]interface{}{
"llm_id": "glm-4-flash@CI@ZHIPU-AI",
"tenant_llm_id": "chat-model-id",
}).Error; err != nil {
t.Fatalf("set default chat model: %v", err)
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
},
},
})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated dataset: %v", err)
}
extractor, ok := persisted.ParserConfig["Extractor:AutoExtractDefault"].(map[string]interface{})
if !ok {
t.Fatalf("extractor parser config = %#v", persisted.ParserConfig["Extractor:AutoExtractDefault"])
}
if extractor["llm_id"] != "chat-model-id" {
t.Fatalf("extractor llm_id = %#v, want tenant model ID", extractor["llm_id"])
}
}
// extractorNodeMetadataInTest returns the metadata object stored on the first
// Extractor node of a parser_config (the component-scoped form).
func extractorNodeMetadataInTest(t *testing.T, cfg map[string]interface{}) map[string]interface{} {
t.Helper()
for cpnID, raw := range cfg {
if !common.IsExtractorComponentID(cpnID) {
continue
}
params, ok := raw.(map[string]interface{})
if !ok {
continue
}
if meta, ok := params["metadata"].(map[string]interface{}); ok {
return meta
}
}
t.Fatalf("no Extractor node with metadata found in parser_config: %#v", cfg)
return nil
}
func TestUpdateDataset_PreservesIncomingMetadataWhenCleaningParserConfig(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
incomingMetadata := []interface{}{map[string]interface{}{
"key": "author",
"type": "string",
"description": "Author",
}}
incomingBuiltInMetadata := []interface{}{map[string]interface{}{
"key": "document_name",
"type": "string",
}}
parserConfig := map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
},
"Extractor:AutoExtractDefault": map[string]interface{}{
"metadata": map[string]interface{}{
"enabled": true,
"metadata": incomingMetadata,
"built_in_metadata": incomingBuiltInMetadata,
},
},
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: parserConfig,
})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err)
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
meta := extractorNodeMetadataInTest(t, map[string]interface{}(persisted.ParserConfig))
if !reflect.DeepEqual(meta, map[string]interface{}{
"enabled": true,
"metadata": incomingMetadata,
"built_in_metadata": incomingBuiltInMetadata,
}) {
t.Fatalf("modular metadata was not preserved on extractor node: %#v", meta)
}
if _, ok := persisted.ParserConfig["metadata"]; ok {
t.Fatalf("top-level flat metadata should be absent: %#v", persisted.ParserConfig["metadata"])
}
if _, ok := persisted.ParserConfig["enable_metadata"]; ok {
t.Fatalf("enable_metadata should be absent: %#v", persisted.ParserConfig["enable_metadata"])
}
if _, ok := persisted.ParserConfig["built_in_metadata"]; ok {
t.Fatalf("built_in_metadata should be absent: %#v", persisted.ParserConfig["built_in_metadata"])
}
if _, ok := persisted.ParserConfig["Parser:HipSignsRhyme"].(map[string]interface{}); !ok {
t.Fatalf("component parser_config missing: %#v", persisted.ParserConfig)
}
}
func TestUpdateDataset_PreservesExistingMetadataWhenParserConfigOmitsIt(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
existingMetadata := []interface{}{map[string]interface{}{
"key": "category",
"type": "string",
}}
if err := dao.DB.Model(&entity.Knowledgebase{}).Where("id = ?", "kb-1").Update("parser_config", entity.JSONMap{
"Extractor:AutoExtractDefault": map[string]interface{}{
"metadata": map[string]interface{}{
"enabled": true,
"metadata": existingMetadata,
"built_in_metadata": []interface{}{map[string]interface{}{"key": "document_name", "type": "string"}},
},
},
}).Error; err != nil {
t.Fatalf("seed parser_config: %v", err)
}
_, code, err := testDatasetUpdateService(t).UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
},
},
})
if err != nil || code != common.CodeSuccess {
t.Fatalf("UpdateDataset failed: code=%d err=%v", code, err)
}
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
meta := extractorNodeMetadataInTest(t, map[string]interface{}(persisted.ParserConfig))
if !reflect.DeepEqual(meta, map[string]interface{}{
"enabled": true,
"metadata": existingMetadata,
"built_in_metadata": []interface{}{map[string]interface{}{"key": "document_name", "type": "string"}},
}) {
t.Fatalf("existing modular metadata was not preserved on extractor node: %#v", meta)
}
if _, ok := persisted.ParserConfig["metadata"]; ok {
t.Fatalf("top-level flat metadata should be absent: %#v", persisted.ParserConfig["metadata"])
}
if _, ok := persisted.ParserConfig["enable_metadata"]; ok {
t.Fatalf("enable_metadata should be absent: %#v", persisted.ParserConfig["enable_metadata"])
}
if _, ok := persisted.ParserConfig["built_in_metadata"]; ok {
t.Fatalf("built_in_metadata should be absent: %#v", persisted.ParserConfig["built_in_metadata"])
}
}
func TestUpdateDataset_StripsCanvasUnknownParam(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
dsl := datasetUpdateCanvasDSL("Parser:CustomRhyme", "pdf")
seedDatasetUpdateCanvas(t, "canvas-1", "tenant-1", dsl)
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:NoSuch": map[string]interface{}{
"pdf": map[string]interface{}{},
},
},
})
if err != nil {
t.Fatalf("UpdateDataset should succeed (unknown cpnID is stripped): %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
persisted, err := dao.NewKnowledgebaseDAO().GetByID(ctx, db, "kb-1")
if err != nil {
t.Fatalf("get updated kb: %v", err)
}
cp := map[string]interface{}(persisted.ParserConfig)
if _, exists := cp["Parser:NoSuch"]; exists {
t.Fatal("expected unknown cpnID Parser:NoSuch to be stripped")
}
// The valid cpnID from the canvas DSL should still be present.
if _, exists := cp["Parser:CustomRhyme"]; !exists {
t.Fatal("expected Parser:CustomRhyme (from canvas DSL defaults) to be present")
}
if result["pipeline_id"] != "canvas-1" {
t.Fatalf("expected pipeline_id preserved, got %#v", result["pipeline_id"])
}
}
func TestUpdateDataset_AcceptsValidCanvasComponentParams(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
dsl := datasetUpdateCanvasDSL("Parser:CustomRhyme", "pdf")
seedDatasetUpdateCanvas(t, "canvas-1", "tenant-1", dsl)
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1")
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
result, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserConfig: map[string]interface{}{
"Parser:CustomRhyme": map[string]interface{}{
"pdf": map[string]interface{}{},
},
},
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
if result["pipeline_id"] != "canvas-1" {
t.Fatalf("expected pipeline_id preserved, got %#v", result["pipeline_id"])
}
}
// TestUpdateDataset_SwitchCanvasToBuiltinValidatesAgainstBuiltin covers the
// effective-ref edge where a request parser_id clears the existing canvas
// pipeline: the override must be validated against the new builtin template,
// not the stale canvas.
func TestUpdateDataset_SwitchCanvasToBuiltinValidatesAgainstBuiltin(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
dsl := datasetUpdateCanvasDSL("Parser:CustomRhyme", "pdf")
seedDatasetUpdateCanvas(t, "canvas-1", "tenant-1", dsl)
insertDatasetUpdateCanvasKB(t, "kb-1", "tenant-1", "Original", "canvas-1")
chunkMethod := "naive"
ctx := withDatasetTestContext(t.Context(), "tenant-1", "tenant-1")
parseType := 1
_, code, err := testDatasetUpdateService(t).UpdateDataset(ctx, "kb-1", service.UpdateDatasetRequest{
ParserID: &chunkMethod,
ParseType: &parseType,
ParserConfig: map[string]interface{}{
// Valid for the "general" builtin template, not the canvas.
"Parser:HipSignsRhyme": map[string]interface{}{
"pdf": map[string]interface{}{"parse_method": "deepdoc"},
},
},
})
if err != nil {
t.Fatalf("UpdateDataset failed: %v", err)
}
if code != common.CodeSuccess {
t.Fatalf("expected success code, got %d", code)
}
}
// TestUpdateDatasetPagerankUnchangedIsANoOpOnInfinity pins Python's guard: the
// engine capability is consulted ONLY when the requested pagerank differs from
// the stored one (`if "pagerank" in req and req["pagerank"] != kb.pagerank`,
// dataset_api_service.py:392). A settings form re-sending the value the dataset
// already has must not fail on a non-ES engine; a real change still must
// (Python's own tests assert that message:
// test/testcases/restful_api/test_datasets.py:620).
func TestUpdateDatasetPagerankUnchangedIsANoOpOnInfinity(t *testing.T) {
db := setupDatasetUpdateTestDB(t)
pushServiceDB(t, db)
insertDatasetUpdateKB(t, "kb-1", "tenant-1", "Original")
// fakeChatDocEngine.SupportsPageRank() == false — the Infinity case.
svc := testDatasetUpdateService(t)
svc.docEngine = fakeChatDocEngine{}
unchanged := int64(0) // insertDatasetUpdateKB stores pagerank 0
if _, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{Pagerank: &unchanged}); err != nil || code != common.CodeSuccess {
t.Fatalf("unchanged pagerank: code=%d err=%v, want success (a no-op, as in Python)", code, err)
}
changed := int64(7)
_, code, err := svc.UpdateDataset(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), "kb-1",
service.UpdateDatasetRequest{Pagerank: &changed})
if err == nil || !strings.Contains(err.Error(), "can only be set when doc_engine is elasticsearch") {
t.Fatalf("changed pagerank: code=%d err=%v, want the doc_engine message", code, err)
}
if code != common.CodeDataError {
t.Fatalf("changed pagerank: code=%d, want CodeDataError", code)
}
persisted, getErr := dao.NewKnowledgebaseDAO().GetByID(withDatasetTestContext(t.Context(), "tenant-1", "tenant-1"), db, "kb-1")
if getErr != nil {
t.Fatalf("get kb: %v", getErr)
}
if persisted.Pagerank != 0 {
t.Fatalf("pagerank = %d, want the rejected update rolled back", persisted.Pagerank)
}
}