notarius/internal/framework/pipeline/registry_integration_test.go

package pipeline

import (
	"context"
	"reflect"
	"testing"

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

func TestRunnerUsesRegistries(t *testing.T) {
	var built []string
	var executed []string
	registries := integrationRegistries(t, &built, &executed)

	output, err := New(registries).Run(context.Background(), RunInput{
		Pipeline: integrationPipeline(),
		SourceID: "source-1",
		RawInput: []byte("source text"),
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}

	wantBuilt := []string{"input", "chunk", "extract-first", "merge", "normalize", "extract-second", "merge", "normalize", "output"}
	if !reflect.DeepEqual(built, wantBuilt) {
		t.Fatalf("built = %#v, want %#v", built, wantBuilt)
	}
	if !reflect.DeepEqual(executed, []string{"extract-first:chunk-0", "extract-second:chunk-0"}) {
		t.Fatalf("executed = %#v, want extractor chunk execution", executed)
	}
	if got := normalizeOutputKeys(output.NormalizeOutputs); !reflect.DeepEqual(got, []string{"normalize", "normalize"}) {
		t.Fatalf("normalize output keys = %#v, want one output from each lane", got)
	}
	if len(output.Rejected) != 0 {
		t.Fatalf("len(Rejected) = %d, want none", len(output.Rejected))
	}
}

func integrationRegistries(t *testing.T, built, executed *[]string) Registries {
	t.Helper()

	registries := Registries{
		Inputs:         NewInputAdapterRegistry(),
		Chunkers:       NewChunkerRegistry(),
		ArtifactCodecs: NewArtifactCodecRegistry(),
		Extractors:     NewExtractorRegistry(),
		Mergers:        NewMergerRegistry(),
		Normalizers:    NewNormalizerRegistry(),
		Outputs:        NewOutputEncoderRegistry(),
	}
	if err := registries.Inputs.Register("input", func() (contracts.InputAdapter, error) {
		*built = append(*built, "input")
		return integrationInput{}, nil
	}); err != nil {
		t.Fatalf("register input: %v", err)
	}
	if err := registries.Chunkers.Register("chunk", func() (contracts.Chunker, error) {
		*built = append(*built, "chunk")
		return integrationChunker{}, nil
	}); err != nil {
		t.Fatalf("register chunker: %v", err)
	}
	registerIntegrationExtractor(t, registries.Extractors, "extract-first", built, executed)
	registerIntegrationExtractor(t, registries.Extractors, "extract-second", built, executed)
	if err := registries.Mergers.Register("merge", func() (contracts.Merger, error) {
		*built = append(*built, "merge")
		return integrationMerger{}, nil
	}); err != nil {
		t.Fatalf("register merger: %v", err)
	}
	if err := registries.Normalizers.Register("normalize", func() (contracts.Normalizer, error) {
		*built = append(*built, "normalize")
		return integrationNormalizer{}, nil
	}); err != nil {
		t.Fatalf("register normalizer: %v", err)
	}
	if err := registries.Outputs.Register("output", func() (contracts.OutputEncoder, error) {
		*built = append(*built, "output")
		return integrationOutput{}, nil
	}); err != nil {
		t.Fatalf("register output: %v", err)
	}
	return registries
}

func registerIntegrationExtractor(t *testing.T, registry *ExtractorRegistry, key string, built, executed *[]string) {
	t.Helper()

	if err := registry.Register(key, func() (contracts.Extractor, error) {
		*built = append(*built, key)
		return integrationExtractor{key: key, executed: executed}, nil
	}); err != nil {
		t.Fatalf("Register(%q) error = %v, want nil", key, err)
	}
}

type integrationInput struct{}

func (input integrationInput) Key() string {
	return "input"
}

func (input integrationInput) Parse(ctx context.Context, req contracts.ParseRequest) (*source.SourceDocument, error) {
	return integrationSourceDocument(), nil
}

type integrationChunker struct{}

func (chunker integrationChunker) Key() string {
	return "chunk"
}

func (chunker integrationChunker) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (chunker integrationChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
	return contracts.ChunkResult{
		Chunks: []source.Chunk{
			{
				ID:       "chunk-0",
				SourceID: req.Source.ID,
				Index:    0,
				Ref: source.SourceRef{
					SourceID:    req.Source.ID,
					StartUnitID: req.Source.Units[0].ID,
					EndUnitID:   req.Source.Units[len(req.Source.Units)-1].ID,
				},
				Content:   []byte(`{"units":[1]}`),
				MediaType: "application/json",
				Units:     req.Source.Units,
			},
		},
	}, nil
}

type integrationExtractor struct {
	key      string
	executed *[]string
}

func (extractor integrationExtractor) Key() string {
	return extractor.key
}

func (extractor integrationExtractor) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (extractor integrationExtractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	*extractor.executed = append(*extractor.executed, extractor.key+":"+req.Chunk.ID)
	return contracts.ExtractionResult{
		Output: contracts.ExtractOutput{
			Schema: contracts.ResponseSchema{ID: "integration", Name: "integration", Version: "v1"},
			Payload: contracts.RawPayload{
				Content:   []byte(`{"value":true}`),
				MediaType: "application/json",
			},
		},
	}, nil
}

type integrationNormalizer struct{}

type integrationMerger struct{}

func (merger integrationMerger) Key() string {
	return "merge"
}

func (merger integrationMerger) Merge(ctx context.Context, req contracts.MergeRequest) (contracts.MergeResult, error) {
	output := contracts.MergeOutput{
		LaneID: req.LaneID,
		Schema: contracts.ResponseSchema{ID: "integration", Name: "integration", Version: "v1"},
		Payload: contracts.RawPayload{
			Content:   []byte(`{"merged":true}`),
			MediaType: "application/json",
		},
	}
	if len(req.ExtractOutputs) > 0 {
		output.SourceID = req.ExtractOutputs[0].SourceID
		output.Schema = req.ExtractOutputs[0].Schema
		output.Payload = req.ExtractOutputs[0].Payload
	}
	return contracts.MergeResult{Output: output}, nil
}

func (normalizer integrationNormalizer) Key() string {
	return "normalize"
}

func (normalizer integrationNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
	return nil
}

func (normalizer integrationNormalizer) Normalize(ctx context.Context, req contracts.NormalizeRequest) (contracts.NormalizeResult, error) {
	return contracts.NormalizeResult{
		Output: contracts.NormalizeOutput{
			LaneID:   req.LaneID,
			SourceID: req.MergeOutput.SourceID,
			Schema:   req.MergeOutput.Schema,
			Payload:  req.MergeOutput.Payload,
		},
	}, nil
}

type integrationOutput struct{}

func (output integrationOutput) Key() string {
	return "output"
}

func (output integrationOutput) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	return contracts.OutputResult{
		Files: []contracts.OutputFile{
			{Name: "output.json", ContentType: "application/json", Bytes: []byte(`{}`)},
		},
	}, nil
}

func integrationPipeline() ResolvedPipeline {
	return ResolvedPipeline{
		ID:     "pipeline-1",
		Digest: "sha256:pipeline",
		Input:  Binding("input"),
		Chunk:  Binding("chunk"),
		ArtifactLanes: []ResolvedArtifactLane{
			{
				ID:        "first",
				Extract:   Binding("extract-first"),
				Merge:     Binding("merge"),
				Normalize: Binding("normalize"),
			},
			{
				ID:        "second",
				Extract:   Binding("extract-second"),
				Merge:     Binding("merge"),
				Normalize: Binding("normalize"),
			},
		},
		Output: Binding("output"),
	}
}

func integrationSourceDocument() *source.SourceDocument {
	return &source.SourceDocument{
		ID:     "source-1",
		Kind:   "document",
		Format: "text/plain",
		Digest: "sha256:abc123",
		Units: []source.SourceUnit{
			{ID: 1, Kind: "unit", Text: "Source unit.", Ref: source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 1}},
		},
	}
}

func normalizeOutputKeys(outputs []contracts.NormalizeOutput) []string {
	keys := make([]string, 0, len(outputs))
	for _, output := range outputs {
		keys = append(keys, output.NormalizerKey)
	}
	return keys
}