92 lines
3.2 KiB
Go
92 lines
3.2 KiB
Go
|
package service
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"errors"
|
||
|
|
"testing"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/models/embedding"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/Tencent/WeKnora/internal/types/interfaces"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
// embeddingInitFailureRepo captures the update made by processChunks. The
|
||
|
|
// real repository has the same UpdateKnowledge contract; this small fake keeps
|
||
|
|
// the test focused on the terminal state transition.
|
||
|
|
type embeddingInitFailureRepo struct {
|
||
|
|
interfaces.KnowledgeRepository
|
||
|
|
knowledge *types.Knowledge
|
||
|
|
updated int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *embeddingInitFailureRepo) GetKnowledgeByID(context.Context, uint64, string) (*types.Knowledge, error) {
|
||
|
|
return r.knowledge, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *embeddingInitFailureRepo) UpdateKnowledge(_ context.Context, knowledge *types.Knowledge) error {
|
||
|
|
r.updated++
|
||
|
|
r.knowledge = knowledge
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type embeddingInitFailureModelService struct {
|
||
|
|
interfaces.ModelService
|
||
|
|
err error
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s embeddingInitFailureModelService) GetEmbeddingModel(context.Context, string) (embedding.Embedder, error) {
|
||
|
|
return nil, s.err
|
||
|
|
}
|
||
|
|
|
||
|
|
// TestProcessChunks_EmbeddingModelFailurePersistsFailedState guards the
|
||
|
|
// failure path that previously just returned from processChunks. A provider
|
||
|
|
// configuration error must be visible to callers and the persisted stage
|
||
|
|
// trace; otherwise the knowledge row remains "processing" indefinitely.
|
||
|
|
func TestProcessChunks_EmbeddingModelFailurePersistsFailedState(t *testing.T) {
|
||
|
|
knowledge := &types.Knowledge{
|
||
|
|
ID: "knowledge-embedding-failure",
|
||
|
|
TenantID: 1,
|
||
|
|
KnowledgeBaseID: "kb-1",
|
||
|
|
ParseStatus: types.ParseStatusProcessing,
|
||
|
|
}
|
||
|
|
repo := &embeddingInitFailureRepo{knowledge: knowledge}
|
||
|
|
failure := errors.New("embedding provider credentials are invalid")
|
||
|
|
|
||
|
|
// Use the real span tracker against the test SQLite schema so this also
|
||
|
|
// proves the embedding stage is failed rather than left running.
|
||
|
|
tracker, db := setupSpanTrackerTest(t)
|
||
|
|
_, attempt, err := tracker.OpenAttempt(context.Background(), knowledge.ID, "trace")
|
||
|
|
require.NoError(t, err)
|
||
|
|
|
||
|
|
tenant := &types.Tenant{ID: 1}
|
||
|
|
ctx := context.WithValue(context.Background(), types.TenantInfoContextKey, tenant)
|
||
|
|
ctx = withAttempt(ctx, attempt)
|
||
|
|
svc := &knowledgeService{
|
||
|
|
repo: repo,
|
||
|
|
modelService: embeddingInitFailureModelService{err: failure},
|
||
|
|
spanTracker: tracker,
|
||
|
|
}
|
||
|
|
kb := &types.KnowledgeBase{
|
||
|
|
ID: "kb-1",
|
||
|
|
EmbeddingModelID: "missing-model",
|
||
|
|
IndexingStrategy: types.IndexingStrategy{VectorEnabled: true},
|
||
|
|
}
|
||
|
|
|
||
|
|
svc.processChunks(ctx, kb, knowledge, nil)
|
||
|
|
|
||
|
|
require.Equal(t, 1, repo.updated, "failure must be persisted immediately")
|
||
|
|
require.Equal(t, types.ParseStatusFailed, knowledge.ParseStatus)
|
||
|
|
require.Equal(t, failure.Error(), knowledge.ErrorMessage)
|
||
|
|
require.False(t, knowledge.UpdatedAt.IsZero())
|
||
|
|
|
||
|
|
var row struct{ Status, ErrorCode, ErrorMessage string }
|
||
|
|
require.NoError(t, db.Table("knowledge_processing_spans").
|
||
|
|
Select("status, error_code, error_message").
|
||
|
|
Where("knowledge_id = ? AND attempt = ? AND name = ?", knowledge.ID, attempt, types.StageEmbedding).
|
||
|
|
First(&row).Error)
|
||
|
|
require.Equal(t, types.SpanStatusFailed, row.Status)
|
||
|
|
require.Equal(t, "EMBEDDING_PROVIDER_FAIL", row.ErrorCode)
|
||
|
|
require.Contains(t, row.ErrorMessage, "get embedding model failed")
|
||
|
|
}
|