story/internal/workflowbridge/service_test.go

package workflowbridge

import (
	"bytes"
	"context"
	"data-terminal-cloud/internal/dify"
	"data-terminal-cloud/internal/extraction"
	"data-terminal-cloud/internal/memoir"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
)

type fixtureEditor struct {
	calls atomic.Int32
	input string
	mu    sync.Mutex
}

func (e *fixtureEditor) GenerateJSON(_ context.Context, input, prompt string) (json.RawMessage, error) {
	e.calls.Add(1)
	raw, _ := json.Marshal(extraction.GenerationResult{Meaningful: true, StoryWorthy: true, DecisionReason: "有原话支持的历史事件", Summary: "进厂往事", Keywords: []string{"进厂"}, Confidence: .9, Stories: []extraction.EventStory{{Title: "进厂", Content: "1970年进厂。", FactContent: "1970年进厂", EventTime: "1970年", SourceIDs: []string{"turn-1"}}}})
	var out map[string]any
	_ = json.Unmarshal(raw, &out)
	out["daily_events"] = []any{}
	return json.Marshal(out)
}
func (e *fixtureEditor) GenerateJSONSchema(_ context.Context, input, prompt string, schema json.RawMessage) (json.RawMessage, error) {
	e.calls.Add(1)
	e.mu.Lock()
	e.input = input
	e.mu.Unlock()
	var in struct {
		Stories []memoir.Source `json:"stories"`
	}
	_ = json.Unmarshal([]byte(input), &in)
	d := memoir.Document{Title: "测试回忆录"}
	for _, s := range in.Stories {
		d.Chapters = append(d.Chapters, memoir.Chapter{StoryID: s.ID, EventTime: s.EventTime, Title: s.Title, Paragraphs: []string{s.Content}, IllustrationPrompt: "水彩情景插画"})
	}
	return json.Marshal(d)
}

type overdueImages struct{}

func (overdueImages) Generate(context.Context, string) ([]byte, error) {
	return nil, extraction.ErrModelAccountOverdue
}

func postStep(t *testing.T, base, step string, ticket dify.Ticket, index int) map[string]any {
	t.Helper()
	raw, _ := json.Marshal(stepInput{ID: ticket.ID, Token: ticket.Token, Index: index})
	r, err := http.Post(base+"/steps/"+step, "application/json", bytes.NewReader(raw))
	if err != nil {
		t.Fatal(err)
	}
	defer r.Body.Close()
	var out map[string]any
	if json.NewDecoder(r.Body).Decode(&out) != nil {
		t.Fatal("invalid step response")
	}
	return out
}

func TestStepsAreCachedAndNeverExposePrivateContent(t *testing.T) {
	editor := &fixtureEditor{}
	s := New("backend-secret", editor, editor, overdueImages{})
	srv := httptest.NewServer(s)
	defer srv.Close()
	b := dify.NewBridge(srv.URL, "backend-secret")
	ticket, err := b.Register(context.Background(), "extract", extraction.GenerationInput{Content: "[turn-1] 长者：1970年进厂。私密原话标记", PreviousStories: []extraction.PreviousStory{}})
	if err != nil {
		t.Fatal(err)
	}
	defer b.Forget(ticket)
	var wg sync.WaitGroup
	for i := 0; i < 5; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			out := postStep(t, srv.URL, "classify", ticket, 0)
			raw, _ := json.Marshal(out)
			if out["success"] != true || strings.Contains(string(raw), "私密原话") || strings.Contains(string(raw), "1970") {
				t.Errorf("private content in metadata %s", raw)
			}
		}()
	}
	wg.Wait()
	if editor.calls.Load() != 1 {
		t.Fatalf("duplicate model calls %d", editor.calls.Load())
	}
	req, _ := http.NewRequest("GET", srv.URL+"/contexts/"+ticket.ID, nil)
	req.Header.Set("Authorization", "Bearer "+ticket.Token)
	resp, _ := http.DefaultClient.Do(req)
	resp.Body.Close()
	if resp.StatusCode != 403 {
		t.Fatal("step token could read private results")
	}
	b.Forget(ticket)
	if s.find(ticket.ID) != nil {
		t.Fatal("deleted context survives")
	}
}

func TestWholeCatalogAndBillingDuringImagesKeepText(t *testing.T) {
	editor := &fixtureEditor{}
	s := New("backend-secret", editor, editor, overdueImages{})
	srv := httptest.NewServer(s)
	defer srv.Close()
	b := dify.NewBridge(srv.URL, "backend-secret")
	sources := []memoir.Source{{ID: "one", Title: "进厂", Content: "1970年进厂。", EventTime: "1970年"}, {ID: "two", Title: "学徒", Content: "1971年学做手套。", EventTime: "1971年"}}
	ticket, err := b.Register(context.Background(), "memoir", map[string]any{"stories": sources})
	if err != nil {
		t.Fatal(err)
	}
	defer b.Forget(ticket)
	if postStep(t, srv.URL, "write-book", ticket, 0)["success"] != true {
		t.Fatal("whole book failed")
	}
	if postStep(t, srv.URL, "validate-book", ticket, 0)["success"] != true {
		t.Fatal("validation failed")
	}
	for i := range sources {
		if postStep(t, srv.URL, "illustrate", ticket, i)["error_code"] != "MODEL_ACCOUNT_OVERDUE" {
			t.Fatal("billing error not identified")
		}
	}
	if postStep(t, srv.URL, "finish", ticket, 0)["success"] != true {
		t.Fatal("image billing failure discarded text")
	}
	var result struct {
		Document memoir.Document `json:"document"`
	}
	if b.Result(context.Background(), ticket, &result) != nil || len(result.Document.Chapters) != 2 || result.Document.IllustrationErrors != 2 {
		t.Fatal("whole text unavailable")
	}
	if editor.calls.Load() != 1 || !strings.Contains(editor.input, "one") || !strings.Contains(editor.input, "two") {
		t.Fatal("model did not receive whole catalog once")
	}
}

func TestPrivateResultPreservesEmptyEventArrays(t *testing.T) {
	s := New("backend-secret", nil, nil, nil)
	srv := httptest.NewServer(s)
	defer srv.Close()
	b := dify.NewBridge(srv.URL, "backend-secret")
	ticket, err := b.Register(context.Background(), "extract", extraction.GenerationInput{Content: "今天吃早饭。"})
	if err != nil {
		t.Fatal(err)
	}
	defer b.Forget(ticket)
	var out struct {
		Result extraction.GenerationResult `json:"result"`
	}
	if err = b.Result(context.Background(), ticket, &out); err != nil || out.Result.Stories == nil || out.Result.DailyEvents == nil {
		t.Fatalf("empty arrays lost in private transport: %v", err)
	}
}