xinmi/internal/application/service/extract_graph_budget_test.go

package service

import (
	"context"
	"crypto/sha256"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"sync"
	"sync/atomic"
	"testing"

	"github.com/Tencent/WeKnora/internal/models/chat"
	"github.com/Tencent/WeKnora/internal/models/requestbudget"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/Tencent/WeKnora/internal/utils"
	"github.com/hibiken/asynq"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"
)

const graphBudgetSource = "Trust supports learning."

func graphBudgetPayload() types.ExtractChunkPayload {
	return types.ExtractChunkPayload{TenantID: 7, ChunkID: "chunk", ModelID: "model",
		ExpectedSource: &types.GraphSourceIdentity{KnowledgeBaseID: "kb", KnowledgeID: "doc", ContentSHA256: fmt.Sprintf("%x", sha256.Sum256([]byte(graphBudgetSource)))}}
}

func openGraphBudgetDB(t *testing.T, path string) *gorm.DB {
	t.Helper()
	db, err := gorm.Open(sqlite.Open(path), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
	require.NoError(t, err)
	sqlDB, err := db.DB()
	require.NoError(t, err)
	sqlDB.SetMaxOpenConns(1)
	t.Cleanup(func() { _ = sqlDB.Close() })
	require.NoError(t, db.Exec(`CREATE TABLE IF NOT EXISTS native_graph_backfill_failures (
		task_id TEXT PRIMARY KEY, tenant_id INTEGER, knowledge_base_id TEXT, chunk_id TEXT,
		knowledge_id TEXT, content_hash TEXT, state TEXT, manual_retries INTEGER, reason TEXT, updated_at DATETIME)`).Error)
	return db
}

func seedGraphBudget(t *testing.T, db *gorm.DB) {
	t.Helper()
	p := graphBudgetPayload()
	require.NoError(t, db.Exec(`INSERT INTO native_graph_backfill_failures VALUES(?,?,?,?,?,?,'dispatch_requested',0,'',CURRENT_TIMESTAMP)`,
		"task", p.TenantID, p.ExpectedSource.KnowledgeBaseID, p.ChunkID, p.ExpectedSource.KnowledgeID, p.ExpectedSource.ContentSHA256).Error)
}

func graphBudgetDB(t *testing.T) *gorm.DB {
	t.Helper()
	db := openGraphBudgetDB(t, filepath.Join(t.TempDir(), "budget.sqlite"))
	seedGraphBudget(t, db)
	return db
}

func TestGraphExecutionClaimSurvivesRestartAndConcurrentDelivery(t *testing.T) {
	path := filepath.Join(t.TempDir(), "budget.sqlite")
	db := openGraphBudgetDB(t, path)
	seedGraphBudget(t, db)
	svc := &ChunkExtractService{executionDB: db}
	var successes atomic.Int32
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if svc.claimGraphExecution(context.Background(), graphBudgetPayload(), "task") == nil {
				successes.Add(1)
			}
		}()
	}
	wg.Wait()
	require.Equal(t, int32(1), successes.Load())
	sqlDB, _ := db.DB()
	require.NoError(t, sqlDB.Close())
	svc = &ChunkExtractService{executionDB: openGraphBudgetDB(t, path)}
	require.ErrorIs(t, svc.claimGraphExecution(context.Background(), graphBudgetPayload(), "task"), requestbudget.ErrStopped)
}

func TestGraphExecutionClaimRejectsChangedOrFailedReceipt(t *testing.T) {
	for _, changed := range []string{"task", "tenant", "base", "chunk", "document", "source", "failed", "retry", "reason", "missing db"} {
		t.Run(changed, func(t *testing.T) {
			db := graphBudgetDB(t)
			p, task := graphBudgetPayload(), "task"
			svc := &ChunkExtractService{executionDB: db}
			switch changed {
			case "task":
				task = "other"
			case "tenant":
				p.TenantID++
			case "base":
				p.ExpectedSource.KnowledgeBaseID = "other"
			case "chunk":
				p.ChunkID = "other"
			case "document":
				p.ExpectedSource.KnowledgeID = "other"
			case "source":
				p.ExpectedSource.ContentSHA256 = "other"
			case "failed":
				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET state='failed'").Error)
			case "retry":
				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET manual_retries=1").Error)
			case "reason":
				require.NoError(t, db.Exec("UPDATE native_graph_backfill_failures SET reason='semantic_review_rejected'").Error)
			case "missing db":
				svc.executionDB = nil
			}
			require.ErrorIs(t, svc.claimGraphExecution(context.Background(), p, task), requestbudget.ErrStopped)
		})
	}
}

func TestGraphExecutionCorrectionSharesActualSendBudget(t *testing.T) {
	t.Setenv("SSRF_WHITELIST", "127.0.0.1")
	utils.ResetSSRFWhitelistForTest()
	t.Cleanup(utils.ResetSSRFWhitelistForTest)
	for _, tc := range []struct {
		name          string
		everyRedirect bool
		wrongModel    bool
	}{
		{name: "three sends complete correction"},
		{name: "fourth send refuses graph write", everyRedirect: true},
		{name: "wrong concrete model sends nothing", wrongModel: true},
	} {
		t.Run(tc.name, func(t *testing.T) {
			var sends, calls atomic.Int32
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				sends.Add(1)
				if r.URL.Path != "/result" {
					call := calls.Add(1)
					if tc.everyRedirect || call == 2 {
						http.Redirect(w, r, "/result", http.StatusTemporaryRedirect)
						return
					}
				}
				content := `[]`
				if calls.Load() == 2 {
					content = `[{"entity1":"Trust","entity2":"learning","relation":"supports","evidence":"Trust supports learning."}]`
				}
				w.Header().Set("Content-Type", "application/json")
				_ = json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"role": "assistant", "content": content}, "finish_reason": "stop"}}})
			}))
			defer server.Close()
			modelID := "model"
			if tc.wrongModel {
				modelID = "other"
			}
			model, err := chat.NewChat(&chat.ChatConfig{Source: types.ModelSourceRemote, BaseURL: server.URL, ModelID: modelID, ModelName: "gpt-4o", APIKey: "local-only", Provider: "openai"}, nil)
			require.NoError(t, err)
			db := graphBudgetDB(t)
			writer := &extractionWriter{}
			svc := &ChunkExtractService{template: &types.PromptTemplateStructured{Description: "Extract graph"}, modelService: extractionModels{model: model},
				knowledgeBaseRepo: extractionKB{}, knowledgeRepo: &extractionKnowledge{}, executionDB: db,
				chunkRepo: extractionChunks{chunk: &types.Chunk{ID: "chunk", TenantID: 7, KnowledgeBaseID: "kb", KnowledgeID: "doc", Content: graphBudgetSource, IsEnabled: true, ChunkType: types.ChunkTypeText}}, graphEngine: writer}
			payload, _ := json.Marshal(graphBudgetPayload())
			task := asynq.NewTask(types.TypeChunkExtract, payload)
			err = svc.handle(context.Background(), task, "task")
			if tc.wrongModel {
				require.Error(t, err)
				require.Zero(t, sends.Load())
				require.Nil(t, writer.saved)
			} else if tc.everyRedirect {
				require.Error(t, err)
				require.Equal(t, int32(3), sends.Load())
				require.Nil(t, writer.saved)
			} else {
				require.NoError(t, err)
				require.Equal(t, int32(3), sends.Load())
				require.True(t, writer.complete)
			}
			if !writer.complete {
				before := sends.Load()
				err = svc.handle(context.Background(), task, "task")
				require.True(t, errors.Is(err, asynq.SkipRetry))
				require.Equal(t, before, sends.Load(), "redelivery must not replenish sends")
			}
		})
	}
}