thread-master/old/backend/internal/library/matrix/generate_test.go

package matrix

import (
	"context"
	"errors"
	"testing"

	domai "haixun-backend/internal/model/ai/domain/usecase"
)

type stubTextGenerator struct {
	responses []string
	errs      []error
	calls     int
}

func (s *stubTextGenerator) GenerateText(_ context.Context, _ domai.GenerateRequest) (*domai.GenerateResult, error) {
	idx := s.calls
	s.calls++
	if idx < len(s.errs) && s.errs[idx] != nil {
		return nil, s.errs[idx]
	}
	if idx >= len(s.responses) {
		return nil, errors.New("no stub response")
	}
	return &domai.GenerateResult{Text: s.responses[idx]}, nil
}

func TestParseGenerateOutputCodeFenceWithPrefix(t *testing.T) {
	raw := "以下是結果：\n```json\n{\"rows\":[{\"sort_order\":1,\"text\":\"hello\",\"angle\":\"a\",\"hook\":\"h\",\"reference_notes\":\"n\",\"source_permalinks\":[],\"rationale\":\"r\"}]}\n```\n"
	got, err := ParseGenerateOutput(raw)
	if err != nil {
		t.Fatalf("ParseGenerateOutput() error = %v", err)
	}
	if len(got.Rows) != 1 || got.Rows[0].Text != "hello" {
		t.Fatalf("rows = %+v", got.Rows)
	}
}

func TestParseGenerateOutputRepairsTruncatedJSON(t *testing.T) {
	raw := `{"rows":[{"sort_order":1,"search_tag":"t","angle":"a","hook":"h","text":"主文","reference_notes":"n","source_permalinks":[],"rationale":"r"},{"sort_order":2,"search_tag":"t2","angle":"b","hook":"h2","text":"第二篇`
	got, err := ParseGenerateOutput(raw)
	if err != nil {
		t.Fatalf("ParseGenerateOutput() error = %v", err)
	}
	if len(got.Rows) != 1 {
		t.Fatalf("rows = %+v, want 1 recovered row", got.Rows)
	}
}

func TestCopyMatrixGenerateRequest_ScalesWithCount(t *testing.T) {
	req := CopyMatrixGenerateRequest(domai.GenerateRequest{}, 5)
	if req.MaxTokens == nil || *req.MaxTokens < 8192 {
		t.Fatalf("max_tokens = %+v, want >= 8192", req.MaxTokens)
	}
}

func TestGenerateCopyOutputRetriesOnParseFailure(t *testing.T) {
	bad := `{"rows":[{"sort_order":1,"text":"`
	good := `{"rows":[{"sort_order":1,"text":"ok","angle":"a","hook":"h","reference_notes":"n","source_permalinks":[],"rationale":"r"}]}`
	gen := &stubTextGenerator{responses: []string{bad, good}}
	got, err := GenerateCopyOutput(context.Background(), gen, domai.GenerateRequest{
		Messages: []domai.Message{{Role: "user", Content: "go"}},
	}, 1)
	if err != nil {
		t.Fatalf("GenerateCopyOutput() error = %v", err)
	}
	if len(got.Rows) != 1 || got.Rows[0].Text != "ok" {
		t.Fatalf("rows = %+v", got.Rows)
	}
	if gen.calls != 2 {
		t.Fatalf("calls = %d, want retry", gen.calls)
	}
}