story/internal/extraction/processor_safety_test.go

package extraction

import (
	"context"
	"testing"

	"data-terminal-cloud/internal/contentpolicy"
)

type safetyTestGenerator struct{ calls int }

func (g *safetyTestGenerator) Generate(context.Context, GenerationInput) (GenerationResult, error) {
	g.calls++
	return GenerationResult{
		Meaningful: true, StoryWorthy: true, DecisionReason: "测试故事",
		Summary: "模型摘要", OpenTopic: "开放话题", StoryTitle: "标题",
		StoryContent: "模型输出风险", FactContent: "普通事实", AllowStory: true,
	}, nil
}

type safetyTestClassifier struct{ calls int }

func (c *safetyTestClassifier) Classify(_ context.Context, _ string) (contentpolicy.Decision, error) {
	c.calls++
	if c.calls == 1 {
		return contentpolicy.Decision{AllowOrdinaryStory: true, Sensitivity: "normal", ProviderRef: "managed-content-safety:test-v1"}, nil
	}
	return contentpolicy.Decision{AllowOrdinaryStory: false, Sensitivity: "sensitive", Reason: "model_output_safety", ProviderRef: "managed-content-safety:test-v1"}, nil
}

func TestGenerateSafelyBlocksModelOutputBeforePersistence(t *testing.T) {
	generator := &safetyTestGenerator{}
	classifier := &safetyTestClassifier{}
	processor := &Processor{generator: generator, safety: classifier}

	result, err := processor.generateSafely(context.Background(), "普通长者话轮")
	if err != nil {
		t.Fatal(err)
	}
	if generator.calls != 1 || classifier.calls != 2 || result.AllowStory || result.ModelRef != "managed-content-safety:test-v1" || result.PromptVersion != "summary-sensitive-filter-v1:model_output_safety" {
		t.Fatalf("result=%+v generator_calls=%d classifier_calls=%d", result, generator.calls, classifier.calls)
	}
}