notarius/internal/framework/pipeline/walking_skeleton_test.go

package pipeline

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"os"
	"reflect"
	"strings"
	"testing"

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

func TestWalkingSkeletonFixture(t *testing.T) {
	inputBytes := readTestFixture(t, "testdata/walking_skeleton_input.json")
	expectedBytes := readTestFixture(t, "testdata/walking_skeleton_output.json")
	llmClient := &walkingSkeletonLLMClient{}

	resolved, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, walkingSkeletonCatalog(t))
	if err != nil {
		t.Fatalf("ResolvePipeline() error = %v, want nil", err)
	}

	output, err := New(walkingSkeletonRegistries(t)).Run(context.Background(), RunInput{
		Pipeline:  resolved,
		SourceID:  "fixture-source",
		Path:      "walking_skeleton_input.json",
		RawInput:  inputBytes,
		LLMClient: llmClient,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}
	if output.ContentType != "application/json" {
		t.Fatalf("ContentType = %q, want application/json", output.ContentType)
	}
	assertStructuralJSONEqual(t, output.EncodedOutput, expectedBytes)
	if llmClient.calls != 2 {
		t.Fatalf("LLM calls = %d, want chunk count 2", llmClient.calls)
	}
}

func TestWalkingSkeletonResolutionRejectsMissingCapability(t *testing.T) {
	catalog := walkingSkeletonCatalog(t)
	catalog.Extractors = NewExtractorRegistry()
	if err := catalog.Extractors.RegisterWithSpec(ModuleSpec{
		Key:      "fake/extract",
		Stage:    StageExtract,
		Requires: []string{"missing"},
		Provides: []string{"fake_artifacts"},
	}, func() (contracts.Extractor, error) {
		return walkingSkeletonExtractor{}, nil
	}); err != nil {
		t.Fatalf("RegisterWithSpec() error = %v, want nil", err)
	}

	_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{}, catalog)
	if err == nil {
		t.Fatal("ResolvePipeline() error = nil, want error")
	}
	if !strings.Contains(err.Error(), "missing") {
		t.Fatalf("ResolvePipeline() error = %q, want missing capability", err.Error())
	}
}

func TestWalkingSkeletonResolutionRejectsUnknownOnlyLane(t *testing.T) {
	_, err := ResolvePipeline(walkingSkeletonProfile(), ResolveOptions{Only: []string{"missing"}}, walkingSkeletonCatalog(t))
	if err == nil {
		t.Fatal("ResolvePipeline() error = nil, want error")
	}
	if !strings.Contains(err.Error(), "missing") || !strings.Contains(err.Error(), "not declared") {
		t.Fatalf("ResolvePipeline() error = %q, want unknown lane error", err.Error())
	}
}

func walkingSkeletonProfile() PipelineProfile {
	return PipelineProfile{
		ID:     "walking-skeleton",
		Input:  Binding("fake/input"),
		Chunk:  Binding("fake/chunk"),
		Output: Binding("json"),
		Artifacts: map[string]ArtifactLaneProfile{
			"events": {Extract: Binding("fake/extract")},
		},
	}
}

func walkingSkeletonCatalog(t *testing.T) ModuleCatalog {
	t.Helper()

	catalog := ModuleCatalog{
		Inputs:      NewInputAdapterRegistry(),
		Chunkers:    NewChunkerRegistry(),
		Extractors:  NewExtractorRegistry(),
		Mergers:     NewMergerRegistry(),
		Normalizers: NewNormalizerRegistry(),
		Outputs:     NewOutputEncoderRegistry(),
	}
	if err := catalog.Inputs.RegisterWithSpec(ModuleSpec{
		Key:      "fake/input",
		Stage:    StageInput,
		Provides: []string{"plain_text"},
	}, func() (contracts.InputAdapter, error) {
		return walkingSkeletonInput{}, nil
	}); err != nil {
		t.Fatalf("register fake input: %v", err)
	}
	if err := catalog.Chunkers.RegisterWithSpec(ModuleSpec{
		Key:      "fake/chunk",
		Stage:    StageChunk,
		Requires: []string{"plain_text"},
		Provides: []string{"chunks"},
	}, func() (contracts.Chunker, error) {
		return walkingSkeletonChunker{}, nil
	}); err != nil {
		t.Fatalf("register fake chunker: %v", err)
	}
	if err := catalog.Extractors.RegisterWithSpec(ModuleSpec{
		Key:      "fake/extract",
		Stage:    StageExtract,
		Requires: []string{"chunks"},
		Provides: []string{"fake_artifacts"},
	}, func() (contracts.Extractor, error) {
		return walkingSkeletonExtractor{}, nil
	}); err != nil {
		t.Fatalf("register fake extractor: %v", err)
	}
	if err := catalog.Mergers.RegisterWithSpec(ModuleSpec{
		Key:      DefaultMergeModule,
		Stage:    StageMerge,
		Requires: []string{"fake_artifacts"},
	}, func() (contracts.Merger, error) {
		return AppendOrderMerger{}, nil
	}); err != nil {
		t.Fatalf("register append-order merger: %v", err)
	}
	if err := catalog.Normalizers.RegisterWithSpec(ModuleSpec{
		Key:   DefaultNormalizeModule,
		Stage: StageNormalize,
	}, func() (contracts.Normalizer, error) {
		return NoopNormalizer{}, nil
	}); err != nil {
		t.Fatalf("register no-op normalizer: %v", err)
	}
	if err := catalog.Outputs.RegisterWithSpec(ModuleSpec{
		Key:   "json",
		Stage: StageOutput,
	}, func() (contracts.OutputEncoder, error) {
		return walkingSkeletonOutput{}, nil
	}); err != nil {
		t.Fatalf("register fake output: %v", err)
	}
	return catalog
}

func walkingSkeletonRegistries(t *testing.T) Registries {
	t.Helper()

	catalog := walkingSkeletonCatalog(t)
	return Registries{
		Inputs:      catalog.Inputs,
		Chunkers:    catalog.Chunkers,
		Extractors:  catalog.Extractors,
		Mergers:     catalog.Mergers,
		Normalizers: catalog.Normalizers,
		Outputs:     catalog.Outputs,
	}
}

type walkingSkeletonInput struct{}

func (input walkingSkeletonInput) Key() string {
	return "fake/input"
}

func (input walkingSkeletonInput) Parse(ctx context.Context, req contracts.ParseRequest) (*source.SourceDocument, error) {
	var fixture struct {
		ID    string `json:"id"`
		Units []struct {
			ID   string `json:"id"`
			Text string `json:"text"`
		} `json:"units"`
	}
	if err := json.Unmarshal(req.Raw, &fixture); err != nil {
		return nil, err
	}

	units := make([]source.SourceUnit, 0, len(fixture.Units))
	for _, unit := range fixture.Units {
		units = append(units, source.SourceUnit{
			ID:   unit.ID,
			Kind: "unit",
			Text: unit.Text,
		})
	}
	return &source.SourceDocument{
		ID:     fixture.ID,
		Kind:   "fixture",
		Format: "application/json",
		Digest: rawDigest(req.Raw),
		Units:  units,
	}, nil
}

type walkingSkeletonChunker struct{}

func (chunker walkingSkeletonChunker) Key() string {
	return "fake/chunk"
}

func (chunker walkingSkeletonChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
	if len(req.Source.Units) < 3 {
		return contracts.ChunkResult{}, fmt.Errorf("fixture source must contain at least three units")
	}
	return contracts.ChunkResult{
		Chunks: []contracts.SourceChunk{
			{
				ID:       req.Source.ID + ":chunk:0",
				SourceID: req.Source.ID,
				Index:    0,
				Units:    append([]source.SourceUnit(nil), req.Source.Units[:2]...),
			},
			{
				ID:       req.Source.ID + ":chunk:1",
				SourceID: req.Source.ID,
				Index:    1,
				Units:    append([]source.SourceUnit(nil), req.Source.Units[2:]...),
			},
		},
	}, nil
}

type walkingSkeletonExtractor struct{}

func (extractor walkingSkeletonExtractor) Key() string {
	return "fake/extract"
}

func (extractor walkingSkeletonExtractor) ArtifactType() string {
	return "fake_event"
}

func (extractor walkingSkeletonExtractor) SchemaVersion() string {
	return "v1"
}

func (extractor walkingSkeletonExtractor) Validators() []contracts.Validator {
	return nil
}

func (extractor walkingSkeletonExtractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	var response struct {
		Call int `json:"call"`
	}
	if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:          "fake/extract",
		ResponseSchemaName: "fake_event",
	}, &response); err != nil {
		return contracts.ExtractionResult{}, err
	}

	payload, err := json.Marshal(map[string]any{
		"chunk_id": req.Chunk.ID,
		"llm_call": response.Call,
		"text":     chunkText(req.Chunk.Units),
	})
	if err != nil {
		return contracts.ExtractionResult{}, err
	}

	return contracts.ExtractionResult{
		Candidates: []artifacts.ArtifactCandidate{
			{
				Payload: payload,
				SourceRefs: []source.SourceRef{
					{
						SourceID:    req.Source.ID,
						StartUnitID: req.Chunk.Units[0].ID,
						EndUnitID:   req.Chunk.Units[len(req.Chunk.Units)-1].ID,
					},
				},
			},
		},
	}, nil
}

type walkingSkeletonLLMClient struct {
	calls int
}

func (client *walkingSkeletonLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
	client.calls++
	if response, ok := out.(*struct {
		Call int `json:"call"`
	}); ok {
		response.Call = client.calls
	}
	content, err := json.Marshal(map[string]any{"call": client.calls})
	if err != nil {
		return contracts.StructuredCompletionResponse{}, err
	}
	return contracts.StructuredCompletionResponse{
		Content: content,
	}, nil
}

type walkingSkeletonOutput struct{}

func (output walkingSkeletonOutput) Key() string {
	return "json"
}

func (output walkingSkeletonOutput) Encode(ctx context.Context, req contracts.OutputRequest) (contracts.OutputResult, error) {
	encoded, err := json.Marshal(struct {
		Manifest artifacts.RunManifest `json:"manifest"`
		Approved []artifacts.Artifact  `json:"approved"`
	}{
		Manifest: artifacts.RunManifest{
			PipelineID:       req.Manifest.PipelineID,
			PipelineDigest:   req.Manifest.PipelineDigest,
			ArtifactLanes:    req.Manifest.ArtifactLanes,
			ValidationStatus: req.Manifest.ValidationStatus,
		},
		Approved: req.Approved,
	})
	if err != nil {
		return contracts.OutputResult{}, err
	}
	return contracts.OutputResult{
		Bytes:       encoded,
		ContentType: "application/json",
	}, nil
}

func readTestFixture(t *testing.T, path string) []byte {
	t.Helper()

	bytes, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read fixture %q: %v", path, err)
	}
	return bytes
}

func assertStructuralJSONEqual(t *testing.T, gotBytes, wantBytes []byte) {
	t.Helper()

	var got any
	if err := json.Unmarshal(gotBytes, &got); err != nil {
		t.Fatalf("unmarshal actual JSON: %v\n%s", err, gotBytes)
	}
	var want any
	if err := json.Unmarshal(wantBytes, &want); err != nil {
		t.Fatalf("unmarshal expected JSON: %v\n%s", err, wantBytes)
	}
	if !reflect.DeepEqual(got, want) {
		gotFormatted, _ := json.MarshalIndent(got, "", "  ")
		wantFormatted, _ := json.MarshalIndent(want, "", "  ")
		t.Fatalf("actual JSON:\n%s\nwant:\n%s", gotFormatted, wantFormatted)
	}
}

func chunkText(units []source.SourceUnit) string {
	parts := make([]string, 0, len(units))
	for _, unit := range units {
		parts = append(parts, unit.Text)
	}
	return strings.Join(parts, " ")
}

func rawDigest(raw []byte) string {
	sum := sha256.Sum256(raw)
	return "sha256:" + hex.EncodeToString(sum[:])
}