47 lines
1.8 KiB
Go
47 lines
1.8 KiB
Go
|
package workbench
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestExistingEvidenceCanBeRevisedWithoutCrossSegmentReuse(t *testing.T) {
|
||
|
|
s := sqliteWorkbench(t)
|
||
|
|
if err := s.db.Exec("CREATE TABLE knowledges (id TEXT PRIMARY KEY, knowledge_base_id TEXT, folder_path TEXT)").Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
if err := s.db.Exec("INSERT INTO knowledges VALUES ('file', 'kb', 'category')").Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
source := &types.SearchResult{ID: "chunk", KnowledgeID: "file", Content: "Source evidence"}
|
||
|
|
ctxFor := func(user, segment, job string) context.Context {
|
||
|
|
return context.WithValue(workbenchContext(7, user, types.TenantRoleContributor), consumptionContextKey{}, consumptionAttribution{"agent", "task", segment, job})
|
||
|
|
}
|
||
|
|
reserve := func(user, segment, job string, want bool) {
|
||
|
|
t.Helper()
|
||
|
|
allowed, err := s.reserveEvidence(ctxFor(user, segment, job), source)
|
||
|
|
if err != nil || allowed != want {
|
||
|
|
t.Fatalf("reserve %s/%s/%s: allowed=%v err=%v", user, segment, job, allowed, err)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
reserve("alice", "first", "original", true)
|
||
|
|
reserve("alice", "first", "original", true)
|
||
|
|
reserve("alice", "second", "concurrent", false)
|
||
|
|
if err := finishConsumption(s.db, &Job{ID: "original", TenantID: 7, OwnerID: "alice"}, true); err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
reserve("alice", "first", "revision", true)
|
||
|
|
reserve("alice", "second", "new", false)
|
||
|
|
reserve("bob", "second", "other-user", true)
|
||
|
|
if err := s.db.Model(&Consumption{}).Where("owner_id=?", "alice").Update("banned", true).Error; err != nil {
|
||
|
|
t.Fatal(err)
|
||
|
|
}
|
||
|
|
reserve("alice", "first", "banned", false)
|
||
|
|
var count int64
|
||
|
|
if err := s.db.Model(&Consumption{}).Count(&count).Error; err != nil || count != 2 {
|
||
|
|
t.Fatalf("expected one record per user, count=%d err=%v", count, err)
|
||
|
|
}
|
||
|
|
}
|