story/internal/memoryindex/processor_integration_test.go

package memoryindex

import (
	"context"
	"errors"
	"os"
	"strings"
	"testing"
	"time"

	"data-terminal-cloud/internal/developmentembedding"
	"data-terminal-cloud/internal/embeddingprovider"
	"data-terminal-cloud/internal/outbox"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
)

type indexLockProbe struct {
	pool      *pgxpool.Pool
	versionID string
}

func (p indexLockProbe) Embed(ctx context.Context, content string) (embeddingprovider.Result, error) {
	updateContext, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()
	if _, err := p.pool.Exec(updateContext, `UPDATE memory_index_entry SET updated_at=updated_at WHERE memory_version_id=$1::uuid`, p.versionID); err != nil {
		return embeddingprovider.Result{}, err
	}
	return embeddingprovider.Result{Vector: developmentembedding.Vector(content), ProviderRef: "integration-lock-probe"}, nil
}

func TestProcessorActivatesAndInvalidatesMemoryIndex(t *testing.T) {
	databaseURL := os.Getenv("TEST_DATABASE_URL")
	if databaseURL == "" {
		t.Skip("TEST_DATABASE_URL is not set")
	}
	ctx := context.Background()
	pool, err := pgxpool.New(ctx, databaseURL)
	if err != nil {
		t.Fatal(err)
	}
	defer pool.Close()
	if _, err := NewProcessor(nil); err == nil {
		t.Fatal("nil database pool was accepted")
	}
	processor, err := NewProcessor(pool)
	if err != nil {
		t.Fatal(err)
	}
	memoryID, memoryPublicID := uuid.NewString(), uuid.NewString()
	versionID, versionPublicID := uuid.NewString(), uuid.NewString()
	defer func() {
		_, _ = pool.Exec(ctx, `DELETE FROM memory_index_entry WHERE memory_id=$1::uuid`, memoryID)
		_, _ = pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=NULL WHERE id=$1::uuid`, memoryID)
		_, _ = pool.Exec(ctx, `DELETE FROM memory_version WHERE memory_id=$1::uuid`, memoryID)
		_, _ = pool.Exec(ctx, `DELETE FROM curated_memory WHERE id=$1::uuid`, memoryID)
	}()
	if _, err := pool.Exec(ctx, `
		INSERT INTO curated_memory(id,public_id,tenant_id,subject_id,status,sensitivity,valid_from)
		VALUES($1::uuid,$2::uuid,'10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401','ACTIVE','normal',now())`, memoryID, memoryPublicID); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `
		INSERT INTO memory_version(id,public_id,memory_id,version_no,content,source_refs,confirmed_by,change_reason,content_digest)
		VALUES($1::uuid,$2::uuid,$3::uuid,1,$4,'[{"type":"integration"}]','10000000-0000-0000-0000-000000000101','合成索引测试',$5)`,
		versionID, versionPublicID, memoryID, "长者在河边小学教书。", strings.Repeat("a", 64)); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `UPDATE curated_memory SET current_version_id=$1::uuid WHERE id=$2::uuid`, versionID, memoryID); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `
		INSERT INTO memory_index_entry(tenant_id,subject_id,memory_id,memory_version_id,index_type,status)
		VALUES('10000000-0000-0000-0000-000000000001','10000000-0000-0000-0000-000000000401',$1::uuid,$2::uuid,'pgvector','PENDING')`, memoryID, versionID); err != nil {
		t.Fatal(err)
	}
	if err := processor.Handle(ctx, outbox.Event{Payload: []byte(`{}`)}); !errors.Is(err, ErrInvalidIndexEvent) {
		t.Fatalf("invalid event error=%v", err)
	}
	missingID := uuid.NewString()
	if err := processor.Handle(ctx, outbox.Event{Payload: []byte(`{"memory_version_id":"` + missingID + `"}`)}); !errors.Is(err, ErrInvalidIndexEvent) {
		t.Fatalf("missing version error=%v", err)
	}
	event := outbox.Event{Payload: []byte(`{"memory_version_id":"` + versionID + `"}`)}
	processor, err = NewProcessor(pool, indexLockProbe{pool: pool, versionID: versionID})
	if err != nil {
		t.Fatal(err)
	}
	if err := processor.Handle(ctx, event); err != nil {
		t.Fatalf("activate index: %v", err)
	}
	var status, providerRef string
	var dimensions int
	if err := pool.QueryRow(ctx, `SELECT status,provider_ref,vector_dims(embedding) FROM memory_index_entry WHERE memory_version_id=$1::uuid`, versionID).Scan(&status, &providerRef, &dimensions); err != nil {
		t.Fatal(err)
	}
	if status != "ACTIVE" || providerRef != "integration-lock-probe" || dimensions != 1536 {
		t.Fatalf("index status=%s provider=%s dimensions=%d", status, providerRef, dimensions)
	}
	if _, err := pool.Exec(ctx, `UPDATE curated_memory SET status='DISABLED' WHERE id=$1::uuid`, memoryID); err != nil {
		t.Fatal(err)
	}
	if _, err := pool.Exec(ctx, `UPDATE memory_index_entry SET status='PENDING' WHERE memory_version_id=$1::uuid`, versionID); err != nil {
		t.Fatal(err)
	}
	if err := processor.Handle(ctx, event); err != nil {
		t.Fatalf("invalidate index: %v", err)
	}
	var embeddingIsNull bool
	if err := pool.QueryRow(ctx, `SELECT status,embedding IS NULL FROM memory_index_entry WHERE memory_version_id=$1::uuid`, versionID).Scan(&status, &embeddingIsNull); err != nil {
		t.Fatal(err)
	}
	if status != "STALE" || !embeddingIsNull {
		t.Fatalf("ineligible index status=%s embedding_null=%v", status, embeddingIsNull)
	}
}