story/internal/memoir/integration_test.go

package memoir

import (
	"context"
	"data-terminal-cloud/internal/extraction"
	"data-terminal-cloud/internal/identity"
	"data-terminal-cloud/internal/outbox"
	"encoding/json"
	"fmt"
	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgxpool"
	"os"
	"strings"
	"testing"
	"time"
)

type testEditor struct{}

func (testEditor) GenerateJSON(_ context.Context, input, prompt string) (json.RawMessage, error) {
	var in struct {
		Stories []Source `json:"stories"`
	}
	_ = json.Unmarshal([]byte(input), &in)
	d := Document{Title: "针脚里的故事"}
	for _, s := range in.Stories {
		d.Chapters = append(d.Chapters, Chapter{StoryID: s.ID, EventTime: s.EventTime, Title: s.Title, Paragraphs: []string{s.Content}, IllustrationPrompt: "水彩场景", Caption: s.Title})
	}
	return json.Marshal(d)
}

type testImage struct{}

func (testImage) Generate(context.Context, string) ([]byte, error) {
	return []byte{0xff, 0xd8, 0xff, 0xe0}, nil
}

type wholeBookEditor struct{ calls, catalogSize int }

func (g *wholeBookEditor) GenerateJSON(ctx context.Context, input, prompt string) (json.RawMessage, error) {
	var in struct {
		Stories []Source `json:"stories"`
	}
	_ = json.Unmarshal([]byte(input), &in)
	g.calls++
	g.catalogSize = len(in.Stories)
	return (testEditor{}).GenerateJSON(ctx, input, prompt)
}

type waitingImage struct{ started, release chan struct{} }

func (g waitingImage) Generate(ctx context.Context, prompt string) ([]byte, error) {
	g.started <- struct{}{}
	select {
	case <-g.release:
		return (testImage{}).Generate(ctx, prompt)
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

type testExtract struct{}

func (testExtract) Generate(_ context.Context, in extraction.GenerationInput) (extraction.GenerationResult, error) {
	id := strings.Split(strings.Split(in.Content, "]")[0], "[")[1]
	result := extraction.GenerationResult{Meaningful: true, StoryWorthy: true, AllowStory: true, Confidence: .9, ModelRef: "test"}
	for i := 0; i < 7; i++ {
		result.Stories = append(result.Stories, extraction.EventStory{Title: fmt.Sprintf("往事%d", i+1), Content: fmt.Sprintf("%d年我经历的事。", 1970+i), FactContent: fmt.Sprintf("%d年往事", 1970+i), EventTime: fmt.Sprintf("%d年", 1970+i), SourceIDs: []string{id}})
	}
	return result, nil
}
func TestMemoirImportGenerateAndInvalidateChangedSources(t *testing.T) {
	testMemoirImportGenerateAndInvalidateChangedSources(t, false)
}

func TestWorkflowMemoirImportGenerateAndInvalidateChangedSources(t *testing.T) {
	testMemoirImportGenerateAndInvalidateChangedSources(t, true)
}

type fixtureBookWorkflow struct {
	editor *wholeBookEditor
	images waitingImage
}

func (w fixtureBookWorkflow) Generate(ctx context.Context, sources []Source, publish func(Document) error) (Document, [][]byte, error) {
	input, _ := json.Marshal(map[string]any{"stories": sources})
	raw, err := w.editor.GenerateJSON(ctx, string(input), MemoirPrompt)
	if err != nil {
		return Document{}, nil, err
	}
	var doc Document
	_ = json.Unmarshal(raw, &doc)
	doc.WorkflowRunID = "native-workflow-run"
	if err = publish(doc); err != nil {
		return Document{}, nil, err
	}
	pictures := make([][]byte, len(doc.Chapters))
	for i, chapter := range doc.Chapters {
		pictures[i], err = w.images.Generate(ctx, chapter.IllustrationPrompt)
		if err != nil {
			return Document{}, nil, err
		}
	}
	return doc, pictures, nil
}

func testMemoirImportGenerateAndInvalidateChangedSources(t *testing.T, workflow bool) {
	if os.Getenv("TEST_DATABASE_URL") == "" {
		t.Skip("TEST_DATABASE_URL not set")
	}
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	pool, e := pgxpool.New(ctx, os.Getenv("TEST_DATABASE_URL"))
	if e != nil {
		t.Fatal(e)
	}
	defer pool.Close()
	id, elder := uuid.NewString(), uuid.NewString()
	tenant := "10000000-0000-0000-0000-000000000001"
	user := "10000000-0000-0000-0000-000000000101"
	exec := func(sql string, args ...any) {
		t.Helper()
		if _, e := pool.Exec(ctx, sql, args...); e != nil {
			t.Fatal(e)
		}
	}
	exec(`INSERT INTO subject_profile(id,public_id,tenant_id,household_id,preferred_name)VALUES($1::uuid,$2::uuid,$3::uuid,'10000000-0000-0000-0000-000000000201','回忆录验收')`, id, elder, tenant)
	exec(`INSERT INTO consent_record(tenant_id,subject_id,consent_type,policy_version,notice_digest,actor_user_id,granted_at) SELECT $1::uuid,$2::uuid,k,'test','test',$3::uuid,now()-interval '1 hour' FROM unnest(ARRAY['transcript_storage','story_extraction'])k`, tenant, id, user)
	t.Cleanup(func() {
		for _, sql := range []string{`DELETE FROM memoir_document WHERE subject_id=$1::uuid`, `DELETE FROM memoir_generation_job WHERE subject_id=$1::uuid`, `DELETE FROM outbox_event WHERE aggregate_id IN(SELECT id FROM memoir_chat_import WHERE subject_id=$1::uuid)`, `DELETE FROM story_generation_event WHERE story_id IN(SELECT id FROM story_draft WHERE subject_id=$1::uuid)`, `DELETE FROM fact_candidate WHERE subject_id=$1::uuid`, `UPDATE story_draft SET current_revision_id=NULL WHERE subject_id=$1::uuid`, `DELETE FROM story_revision WHERE story_id IN(SELECT id FROM story_draft WHERE subject_id=$1::uuid)`, `DELETE FROM story_draft WHERE subject_id=$1::uuid`, `DELETE FROM memoir_import_turn WHERE import_id IN(SELECT id FROM memoir_chat_import WHERE subject_id=$1::uuid)`, `DELETE FROM memoir_chat_import WHERE subject_id=$1::uuid`, `DELETE FROM consent_record WHERE subject_id=$1::uuid`, `DELETE FROM subject_profile WHERE id=$1::uuid`} {
			_, _ = pool.Exec(ctx, sql, id)
		}
	})
	p := identity.Principal{Issuer: "http://localhost:58180/realms/voice-local", Subject: "00000000-0000-0000-0000-000000000101"}
	editor := &wholeBookEditor{}
	images := waitingImage{started: make(chan struct{}, 7), release: make(chan struct{})}
	s := &Service{Pool: pool, Generator: editor, Illustrator: images, Model: "test"}
	if workflow {
		s.Generator = nil
		s.Workflow = fixtureBookWorkflow{editor, images}
	}
	importID, e := s.Import(ctx, p, ImportInput{ID: uuid.NewString(), ElderID: elder, Transcript: "我：1970年我进厂。\n豆包：你开过飞船？", Title: "授权的测试聊天"})
	if e != nil {
		t.Fatal(e)
	}
	processor, _ := extraction.NewProcessor(pool, testExtract{})
	if e = processor.HandleImport(ctx, outbox.Event{AggregateID: importID}); e != nil {
		t.Fatal(e)
	}
	if _, e = s.Query(ctx, identity.Principal{Issuer: p.Issuer, Subject: "00000000-0000-0000-0000-000000000103"}, elder); e == nil {
		t.Fatal("unrelated account could read memoir")
	}
	job, e := s.Enqueue(ctx, p, elder)
	if e != nil {
		t.Fatal(e)
	}
	finished := make(chan error, 1)
	go func() { finished <- s.Process(ctx, job) }()
	select {
	case <-images.started:
	case e := <-finished:
		t.Fatalf("generation ended before images: %v", e)
	case <-time.After(10 * time.Second):
		t.Fatal("generation did not start images")
	}
	partial, e := s.Query(ctx, p, elder)
	if e != nil || partial["document"] == nil || partial["generation_state"] != "RUNNING" {
		t.Fatalf("text should be readable while images are pending: %v %v", partial, e)
	}
	close(images.release)
	if e = <-finished; e != nil {
		t.Fatal(e)
	}
	if editor.calls != 1 || editor.catalogSize != 7 {
		t.Fatalf("expected one whole-catalog request, calls=%d size=%d", editor.calls, editor.catalogSize)
	}
	data, e := s.Query(ctx, p, elder)
	if e != nil || data["document"] == nil {
		t.Fatalf("memoir unavailable %v", e)
	}
	var book Document
	_ = json.Unmarshal(data["document"].(json.RawMessage), &book)
	if len(book.Chapters) != 7 || book.Chapters[0].ImageID == "" {
		t.Fatal("chapter or illustration missing")
	}
	if workflow && book.WorkflowRunID != "native-workflow-run" {
		t.Fatal("workflow run provenance missing")
	}
	exec(`UPDATE story_draft SET version=version+1 WHERE subject_id=$1::uuid`, id)
	data, e = s.Query(ctx, p, elder)
	if e != nil || data["document"] != nil {
		t.Fatal("old memoir exposed after source version changed")
	}
	exec(`UPDATE consent_record SET withdrawn_at=now() WHERE subject_id=$1::uuid AND consent_type='story_extraction'`, id)
	if _, e = s.Query(ctx, p, elder); e == nil {
		t.Fatal("withdrawn consent still permits memoir")
	}
}