xinmi/internal/workbench/consumption_test.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)
	}
}