notarius/internal/modules/seriatim/input/transcript/runner_helpers_test.go

package transcript

import (
	"context"
	"encoding/json"
	"reflect"
	"testing"

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

func loadPipelineConfig(t *testing.T) config.Config {
	t.Helper()

	cfg := config.Default()
	cfg.Pipelines["seriatim-fixture"] = pipeline.PipelineProfile{
		ID:    "seriatim-fixture",
		Input: pipeline.Binding(Key),
		Chunk: pipeline.Binding("fake/chunk"),
		Artifacts: map[string]pipeline.ArtifactLaneProfile{
			"events": {Extract: pipeline.Binding("fake/extract")},
		},
	}
	return cfg
}

func seriatimTestCatalog(t *testing.T, inputSpec pipeline.ModuleSpec) pipeline.ModuleCatalog {
	t.Helper()

	inputs := pipeline.NewInputAdapterRegistry()
	chunkers := pipeline.NewChunkerRegistry()
	extractors := pipeline.NewExtractorRegistry()
	mergers := pipeline.NewMergerRegistry()
	normalizers := pipeline.NewNormalizerRegistry()
	outputs := pipeline.NewOutputEncoderRegistry()

	if reflect.DeepEqual(inputSpec, ModuleSpec()) {
		if err := Register(inputs); err != nil {
			t.Fatalf("register seriatim input: %v", err)
		}
	} else if err := inputs.RegisterWithSpec(inputSpec, func() (contracts.InputAdapter, error) {
		return New(), nil
	}); err != nil {
		t.Fatalf("register seriatim input override: %v", err)
	}

	mustRegisterChunker(t, chunkers, pipeline.ModuleSpec{Key: "fake/chunk", Stage: pipeline.StageChunk, ExecutionClass: contracts.ExecutionClassDeterministic, Requires: []string{"source.transcript"}, Provides: []string{"chunks"}})
	mustRegisterExtractor(t, extractors, pipeline.ModuleSpec{
		Key: "fake/extract", Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassDeterministic, ArtifactKind: seriatimArtifactKind,
		Requires: []string{"chunks", "transcript.speaker", "transcript.timestamps"}, Provides: []string{"fake.artifacts"},
	})
	mustRegisterMerger(t, mergers, pipeline.ModuleSpec{Key: pipeline.DefaultMergeModule, Stage: pipeline.StageMerge, ExecutionClass: contracts.ExecutionClassDeterministic, ArtifactKind: seriatimArtifactKind, Requires: []string{"fake.artifacts"}})
	mustRegisterNormalizer(t, normalizers, pipeline.ModuleSpec{Key: pipeline.DefaultNormalizeModule, Stage: pipeline.StageNormalize, ExecutionClass: contracts.ExecutionClassDeterministic, ArtifactKind: seriatimArtifactKind})
	mustRegisterOutput(t, outputs, pipeline.ModuleSpec{Key: pipeline.DefaultOutputModule, Stage: pipeline.StageOutput, ExecutionClass: contracts.ExecutionClassDeterministic})

	codecs := pipeline.NewArtifactCodecRegistry()
	if err := pipeline.RegisterArtifactCodec(codecs, seriatimArtifactCodec{}); err != nil {
		t.Fatalf("register artifact codec: %v", err)
	}
	return pipeline.ModuleCatalog{Inputs: inputs, Chunkers: chunkers, ArtifactCodecs: codecs, Extractors: extractors, Mergers: mergers, Normalizers: normalizers, ValidatorChains: pipeline.NewValidatorChainRegistry(), Outputs: outputs}
}

func mustRegisterChunker(t *testing.T, registry *pipeline.ChunkerRegistry, spec pipeline.ModuleSpec) {
	t.Helper()
	if err := registry.RegisterWithSpec(spec, func() (contracts.Chunker, error) { return fakeChunker{}, nil }); err != nil {
		t.Fatalf("register chunker: %v", err)
	}
}

func mustRegisterExtractor(t *testing.T, registry *pipeline.ExtractorRegistry, spec pipeline.ModuleSpec) {
	t.Helper()
	if err := pipeline.RegisterExtractor[seriatimArtifact](registry, spec, func() (contracts.Extractor[seriatimArtifact], error) { return fakeExtractor{}, nil }); err != nil {
		t.Fatalf("register extractor: %v", err)
	}
}

func mustRegisterMerger(t *testing.T, registry *pipeline.MergerRegistry, spec pipeline.ModuleSpec) {
	t.Helper()
	if err := pipeline.RegisterMerger[seriatimArtifact](registry, spec, func() (contracts.Merger[seriatimArtifact], error) { return fakeMerger{}, nil }); err != nil {
		t.Fatalf("register merger: %v", err)
	}
}

func mustRegisterNormalizer(t *testing.T, registry *pipeline.NormalizerRegistry, spec pipeline.ModuleSpec) {
	t.Helper()
	if err := pipeline.RegisterNormalizer[seriatimArtifact](registry, spec, func() (contracts.Normalizer[seriatimArtifact], error) { return fakeNormalizer{}, nil }); err != nil {
		t.Fatalf("register normalizer: %v", err)
	}
}

func mustRegisterOutput(t *testing.T, registry *pipeline.OutputEncoderRegistry, spec pipeline.ModuleSpec) {
	t.Helper()
	if err := registry.RegisterWithSpec(spec, func() (contracts.OutputEncoder, error) { return fakeOutput{}, nil }); err != nil {
		t.Fatalf("register output: %v", err)
	}
}

type fakeChunker struct{}

func (fakeChunker) Key() string                               { return "fake/chunk" }
func (fakeChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
func (fakeChunker) Plan(context.Context, contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
	return contracts.ChunkPlanResult{}, nil
}

type fakeExtractor struct{}

func (fakeExtractor) Key() string                               { return "fake/extract" }
func (fakeExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil }
func (fakeExtractor) Extract(context.Context, contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[seriatimArtifact], error) {
	return contracts.TypedExtractionResult[seriatimArtifact]{}, nil
}

type fakeMerger struct{}

func (fakeMerger) Key() string { return pipeline.DefaultMergeModule }
func (fakeMerger) Merge(_ context.Context, req contracts.TypedMergeRequest[seriatimArtifact]) (contracts.TypedMergeResult[seriatimArtifact], error) {
	if len(req.ExtractOutputs) == 0 {
		return contracts.TypedMergeResult[seriatimArtifact]{}, nil
	}
	return contracts.TypedMergeResult[seriatimArtifact]{Value: req.ExtractOutputs[0].Value}, nil
}

type fakeNormalizer struct{}

func (fakeNormalizer) Key() string                               { return pipeline.DefaultNormalizeModule }
func (fakeNormalizer) ReferenceSlots() []contracts.ReferenceSlot { return nil }
func (fakeNormalizer) Normalize(_ context.Context, req contracts.TypedNormalizeRequest[seriatimArtifact]) (contracts.TypedNormalizeResult[seriatimArtifact], error) {
	return contracts.TypedNormalizeResult[seriatimArtifact]{Value: req.MergeOutput.Value}, nil
}

type fakeOutput struct{}

func (fakeOutput) Key() string { return pipeline.DefaultOutputModule }
func (fakeOutput) Encode(context.Context, contracts.OutputRequest) (contracts.OutputResult, error) {
	return contracts.OutputResult{}, nil
}

const seriatimArtifactKind contracts.ArtifactKind = "test/seriatim-event"

type seriatimArtifact struct {
	Value      string             `json:"value"`
	SourceRefs []source.SourceRef `json:"source_refs"`
}

type seriatimArtifactCodec struct{}

func (seriatimArtifactCodec) Kind() contracts.ArtifactKind { return seriatimArtifactKind }
func (seriatimArtifactCodec) Schema() contracts.ArtifactSchema {
	return contracts.ArtifactSchema{ID: "fake.event", Name: "fake_event", Version: "v1", JSONSchema: []byte(`{"type":"object"}`)}
}
func (seriatimArtifactCodec) MediaType() string { return "application/json" }
func (seriatimArtifactCodec) EncodeCandidate(value seriatimArtifact) ([]byte, error) {
	return json.Marshal(value)
}
func (seriatimArtifactCodec) Encode(value seriatimArtifact) ([]byte, error) {
	return json.Marshal(value)
}
func (seriatimArtifactCodec) Decode(content []byte) (seriatimArtifact, error) {
	var value seriatimArtifact
	err := json.Unmarshal(content, &value)
	return value, err
}