notarius/internal/modules/dnd/register/register_test.go

package register

import (
	"io/fs"
	"reflect"
	"sort"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
)

func TestRegisterAddsDNDFamily(t *testing.T) {
	registries := completeRegistries()
	assets := llm.NewAssetRegistry()
	if err := Register(registries, assets); err != nil {
		t.Fatalf("Register() error = %v, want nil", err)
	}
	assertKeys(t, "chunkers", registries.Chunkers.RegisteredKeys(), []string{"dnd/scenes"})
	assertKeys(t, "extractors", registries.Extractors.RegisteredKeys(), []string{"dnd/spells"})
	assertKeys(t, "validators", registries.Validators.RegisteredKeys(), []string{
		"extract/dnd/spells/shape",
		"extract/dnd/spells/source_refs",
		"extract/dnd/spells/source_relatedness",
	})
	wantChain := []pipeline.ModuleBinding{
		pipeline.Binding("generic/valid_json"),
		pipeline.Binding("generic/valid_json_schema"),
		pipeline.Binding("extract/dnd/spells/shape"),
		pipeline.Binding("extract/dnd/spells/source_refs"),
		pipeline.Binding("extract/dnd/spells/source_relatedness"),
	}
	if got := registries.ValidatorChains.Validators(pipeline.StageExtract, spells.Key); !reflect.DeepEqual(got, wantChain) {
		t.Fatalf("spell validator chain = %#v, want %#v", got, wantChain)
	}
	assertAssetNames(t, assets.PromptFS, []string{
		"dnd.scenes/dnd.scenes.yaml",
		"dnd.scenes/instructions.md",
		"dnd.scenes/sharedassets/common-dnd-references.md",
		"dnd.scenes/sharedassets/common-dnd-system.md",
		"dnd.scenes/sharedassets/common-dnd-transcript.md",
		"dnd.scenes/task.md",
		"dnd.spells/dnd.spells.yaml",
		"dnd.spells/instructions.md",
		"dnd.spells/sharedassets/common-dnd-references.md",
		"dnd.spells/sharedassets/common-dnd-system.md",
		"dnd.spells/sharedassets/common-dnd-transcript.md",
		"dnd.spells/task.md",
	})
	assertAssetNames(t, assets.SchemaFS, []string{
		"dnd_scenes.v1.json",
		"dnd_spells.v1.json",
		"dnd_spells_llm.v1.json",
	})
}

func TestRegisterRejectsMissingDNDDependenciesBeforeMutation(t *testing.T) {
	tests := []struct {
		name    string
		remove  func(*pipeline.Registries, **llm.AssetRegistry)
		wantErr string
	}{
		{name: "chunkers", remove: func(r *pipeline.Registries, _ **llm.AssetRegistry) { r.Chunkers = nil }, wantErr: "chunker registry"},
		{name: "extractors", remove: func(r *pipeline.Registries, _ **llm.AssetRegistry) { r.Extractors = nil }, wantErr: "extractor registry"},
		{name: "validators", remove: func(r *pipeline.Registries, _ **llm.AssetRegistry) { r.Validators = nil }, wantErr: "validator registry"},
		{name: "validator chains", remove: func(r *pipeline.Registries, _ **llm.AssetRegistry) { r.ValidatorChains = nil }, wantErr: "validator chain registry"},
		{name: "assets", remove: func(_ *pipeline.Registries, assets **llm.AssetRegistry) { *assets = nil }, wantErr: "asset registry"},
	}
	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			registries := completeRegistries()
			assets := llm.NewAssetRegistry()
			test.remove(&registries, &assets)
			err := Register(registries, assets)
			if err == nil || !strings.Contains(err.Error(), test.wantErr) {
				t.Fatalf("Register() error = %v, want %q", err, test.wantErr)
			}
			if got := registries.Chunkers; got != nil && len(got.RegisteredKeys()) != 0 {
				t.Fatalf("chunker keys = %#v, want validation before mutation", got.RegisteredKeys())
			}
		})
	}
}

func TestRegisterReportsDuplicateDNDRegistration(t *testing.T) {
	registries := completeRegistries()
	assets := llm.NewAssetRegistry()
	if err := Register(registries, assets); err != nil {
		t.Fatalf("first Register() error = %v, want nil", err)
	}
	err := Register(registries, assets)
	if err == nil || !strings.Contains(err.Error(), "register dnd scenes chunker") || !strings.Contains(err.Error(), "already registered") {
		t.Fatalf("second Register() error = %v, want contextual duplicate error", err)
	}
}

func completeRegistries() pipeline.Registries {
	return pipeline.Registries{
		Inputs:          pipeline.NewInputAdapterRegistry(),
		Chunkers:        pipeline.NewChunkerRegistry(),
		ArtifactCodecs:  pipeline.NewArtifactCodecRegistry(),
		Extractors:      pipeline.NewExtractorRegistry(),
		Mergers:         pipeline.NewMergerRegistry(),
		Normalizers:     pipeline.NewNormalizerRegistry(),
		Validators:      pipeline.NewValidatorRegistry(),
		ValidatorChains: pipeline.NewValidatorChainRegistry(),
		Outputs:         pipeline.NewOutputEncoderRegistry(),
	}
}

func assertKeys(t *testing.T, name string, got, want []string) {
	t.Helper()
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("%s keys = %#v, want %#v", name, got, want)
	}
}

func assertAssetNames(t *testing.T, getFS func() (fs.FS, error), want []string) {
	t.Helper()
	fSys, err := getFS()
	if err != nil {
		t.Fatalf("asset filesystem error = %v, want nil", err)
	}
	var got []string
	if err := fs.WalkDir(fSys, ".", func(path string, entry fs.DirEntry, err error) error {
		if err == nil && !entry.IsDir() {
			got = append(got, path)
		}
		return err
	}); err != nil {
		t.Fatalf("walk assets: %v", err)
	}
	sort.Strings(got)
	if !reflect.DeepEqual(got, want) {
		t.Fatalf("asset names = %#v, want %#v", got, want)
	}
}