xinmi/internal/application/service/knowledge_process_embedding_failure_test.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")
}