notarius/internal/framework/prompt/render_test.go

package prompt

import (
	"crypto/sha256"
	"encoding/hex"
	"strings"
	"testing"
	"testing/fstest"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func TestRenderUserSystemReturnsTextAndMetadata(t *testing.T) {
	system, user, metadata, err := RenderUserSystem(TestGenericPromptID, map[string]any{
		"Task":  "Summarize",
		"Input": "Example input",
	})
	if err != nil {
		t.Fatalf("RenderUserSystem: %v", err)
	}

	if !strings.Contains(system, "generic Notarius test prompt") {
		t.Fatalf("unexpected system prompt: %q", system)
	}
	if !strings.Contains(user, "Task: Summarize") || !strings.Contains(user, "Example input") {
		t.Fatalf("unexpected user prompt: %q", user)
	}
	if strings.TrimSpace(system) != system {
		t.Fatalf("expected trimmed system prompt: %q", system)
	}
	if strings.TrimSpace(user) != user {
		t.Fatalf("expected trimmed user prompt: %q", user)
	}
	if metadata.PromptID != TestGenericPromptID {
		t.Fatalf("unexpected metadata: %+v", metadata)
	}
}

func TestRenderUserSystemUnknownPromptReturnsError(t *testing.T) {
	_, _, _, err := RenderUserSystem("unknown", map[string]any{})
	if err == nil || !strings.Contains(err.Error(), "unknown prompt id") {
		t.Fatalf("expected unknown prompt error, got %v", err)
	}
}

func TestRenderUserSystemMissingTemplateDataReturnsError(t *testing.T) {
	_, _, _, err := RenderUserSystem(TestGenericPromptID, map[string]any{
		"Task": "Summarize",
	})
	if err == nil || !strings.Contains(err.Error(), "Input") {
		t.Fatalf("expected missing template data error, got %v", err)
	}
}

func TestRenderUserSystemIncludesHardeningText(t *testing.T) {
	system, _, _, err := RenderUserSystem(TestGenericPromptID, map[string]any{
		"Task":  "Summarize",
		"Input": "Example input",
	})
	if err != nil {
		t.Fatalf("RenderUserSystem: %v", err)
	}

	hardening := strings.TrimSpace(HardeningText())
	if hardening == "" {
		t.Fatalf("expected hardening text")
	}
	if !strings.Contains(system, hardening) {
		t.Fatalf("expected rendered system prompt to include hardening text: %q", system)
	}
}

func TestRenderUserSystemWithReferencesRendersDeclaredSlots(t *testing.T) {
	bundle := loadReferenceBundle(t,
		[]contracts.ReferenceSlot{{Name: "roster"}, {Name: "glossary"}},
		`System has roster={{ hasreference "roster" }} has glossary={{ hasreference "glossary" }}`,
		`Roster={{ reference "roster" }} Glossary={{ reference "glossary" }}`,
	)
	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"roster": {
			Slot: contracts.ReferenceSlot{Name: "roster"},
			Items: []contracts.ReferenceItem{
				{SlotName: "roster", Content: []byte("Aria")},
			},
		},
	}}

	system, user, _, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err != nil {
		t.Fatalf("RenderUserSystemWithReferences: %v", err)
	}
	if !strings.Contains(system, "has roster=true") || !strings.Contains(system, "has glossary=false") {
		t.Fatalf("system = %q, want reference presence flags", system)
	}
	if !strings.Contains(user, "Roster=Aria") || !strings.Contains(user, "Glossary=") {
		t.Fatalf("user = %q, want rendered and empty optional references", user)
	}
}

func TestRenderUserSystemReferenceHasReferenceRequiresContent(t *testing.T) {
	bundle := loadReferenceBundle(t,
		[]contracts.ReferenceSlot{{Name: "roster"}},
		`System`,
		`{{ hasreference "roster" }} {{ reference "roster" }}`,
	)
	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"roster": {
			Slot: contracts.ReferenceSlot{Name: "roster"},
			Items: []contracts.ReferenceItem{
				{SlotName: "roster", Content: nil},
			},
		},
	}}

	_, user, _, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err != nil {
		t.Fatalf("RenderUserSystemWithReferences: %v", err)
	}
	if user != "false" {
		t.Fatalf("user = %q, want false with empty reference content", user)
	}
}

func TestLoadBundleRejectsUndeclaredReferenceSlots(t *testing.T) {
	_, err := LoadBundle(referenceBundleFS(`System`, `{{ reference "roster" }}`), referenceBundleDefinition(nil))
	if err == nil || !strings.Contains(err.Error(), "roster") || !strings.Contains(err.Error(), "not declared") {
		t.Fatalf("LoadBundle() error = %v, want undeclared reference slot error", err)
	}
}

func TestLoadBundleRejectsDynamicReferenceSlotNames(t *testing.T) {
	_, err := LoadBundle(referenceBundleFS(`System`, `{{ reference .SlotName }}`), referenceBundleDefinition([]contracts.ReferenceSlot{{Name: "roster"}}))
	if err == nil || !strings.Contains(err.Error(), "string literal") {
		t.Fatalf("LoadBundle() error = %v, want string literal error", err)
	}
}

func TestLoadBundleRejectsNestedUndeclaredReferenceSlots(t *testing.T) {
	_, err := LoadBundle(referenceBundleFS(`System`, `{{ printf "%s" (reference "roster") }}`), referenceBundleDefinition(nil))
	if err == nil || !strings.Contains(err.Error(), "roster") || !strings.Contains(err.Error(), "not declared") {
		t.Fatalf("LoadBundle() error = %v, want nested undeclared reference slot error", err)
	}
}

func TestRenderUserSystemRejectsMultipleReferenceItemsUnlessDeclared(t *testing.T) {
	bundle := loadReferenceBundle(t,
		[]contracts.ReferenceSlot{{Name: "roster"}},
		`System`,
		`{{ reference "roster" }}`,
	)
	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"roster": {
			Slot: contracts.ReferenceSlot{Name: "roster"},
			Items: []contracts.ReferenceItem{
				{SlotName: "roster", Content: []byte("Aria")},
				{SlotName: "roster", Content: []byte("Bryn")},
			},
		},
	}}

	_, _, _, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err == nil || !strings.Contains(err.Error(), "does not allow multiple") {
		t.Fatalf("RenderUserSystemWithReferences() error = %v, want multiple item error", err)
	}
}

func TestRenderUserSystemRendersMultipleReferenceItemsDeterministicallyWhenDeclared(t *testing.T) {
	bundle := loadReferenceBundle(t,
		[]contracts.ReferenceSlot{{Name: "roster", Multiple: true}},
		`System`,
		`{{ reference "roster" }}`,
	)
	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"roster": {
			Slot: contracts.ReferenceSlot{Name: "roster", Multiple: true},
			Items: []contracts.ReferenceItem{
				{SlotName: "roster", Content: []byte("Aria")},
				{SlotName: "roster", Content: []byte("Bryn")},
			},
		},
	}}

	_, first, _, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err != nil {
		t.Fatalf("RenderUserSystemWithReferences(first): %v", err)
	}
	_, second, _, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err != nil {
		t.Fatalf("RenderUserSystemWithReferences(second): %v", err)
	}
	if first != "Aria\nBryn" || first != second {
		t.Fatalf("rendered references = %q/%q, want deterministic item order", first, second)
	}
}

func TestPromptMetadataHashIgnoresRenderedReferenceContent(t *testing.T) {
	systemSource := `System`
	userSource := `{{ reference "roster" }}`
	bundle := loadReferenceBundle(t, []contracts.ReferenceSlot{{Name: "roster"}}, systemSource, userSource)
	references := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"roster": {
			Slot:  contracts.ReferenceSlot{Name: "roster"},
			Items: []contracts.ReferenceItem{{SlotName: "roster", Content: []byte("Aria")}},
		},
	}}

	_, _, metadata, err := bundle.RenderUserSystemWithReferences(map[string]any{}, references)
	if err != nil {
		t.Fatalf("RenderUserSystemWithReferences: %v", err)
	}
	hash := sha256.Sum256([]byte(systemSource + "\n\n" + userSource))
	want := "sha256:" + hex.EncodeToString(hash[:])
	if metadata.SHA256 != want {
		t.Fatalf("metadata.SHA256 = %q, want template source hash %q", metadata.SHA256, want)
	}
}

func loadReferenceBundle(t *testing.T, slots []contracts.ReferenceSlot, systemSource string, userSource string) *Bundle {
	t.Helper()
	bundle, err := LoadBundle(referenceBundleFS(systemSource, userSource), referenceBundleDefinition(slots))
	if err != nil {
		t.Fatalf("LoadBundle() error = %v, want nil", err)
	}
	return bundle
}

func referenceBundleDefinition(slots []contracts.ReferenceSlot) Definition {
	return Definition{
		PromptID:       "test.references",
		Version:        VersionV1,
		EmbeddedPath:   "assets/test/references",
		SystemPath:     "assets/test/references/system.md",
		UserPath:       "assets/test/references/user.md",
		ReferenceSlots: slots,
	}
}

func referenceBundleFS(systemSource string, userSource string) fstest.MapFS {
	return fstest.MapFS{
		"assets/test/references/system.md": {Data: []byte(systemSource)},
		"assets/test/references/user.md":   {Data: []byte(userSource)},
	}
}