notarius/internal/modules/dnd/shared/extraction_inputs_test.go

package shared

import (
	"context"
	"reflect"
	"strings"
	"testing"

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

func TestPrepareChunkExtraction(t *testing.T) {
	canceled, cancel := context.WithCancel(context.Background())
	cancel()
	tests := []struct {
		name         string
		ctx          context.Context
		configure    func(*contracts.TypedExtractionRequest)
		sourceInput  contracts.LLMInputMaterial
		want         contracts.LLMInputMaterial
		wantErr      string
		mutateOutput bool
	}{
		{
			name:    "nil context",
			wantErr: "context must not be nil",
		},
		{
			name:    "canceled context",
			ctx:     canceled,
			wantErr: "context error before extraction",
		},
		{
			name: "nil source",
			configure: func(req *contracts.TypedExtractionRequest) {
				req.Source = nil
			},
			wantErr: "source must not be nil",
		},
		{
			name: "nil chunk",
			configure: func(req *contracts.TypedExtractionRequest) {
				req.Chunk = nil
			},
			wantErr: "chunk must not be nil",
		},
		{
			name: "empty chunk units",
			configure: func(req *contracts.TypedExtractionRequest) {
				req.Chunk = &source.Chunk{ID: req.Chunk.ID}
			},
			wantErr: "units must not be empty",
		},
		{
			name:         "fallback to chunk content",
			want:         contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"units":[1,2]}`), "", ""),
			mutateOutput: true,
		},
		{
			name:         "clone isolation",
			sourceInput:  contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"units":[1,2]}`), "sha256:source", "file:///source.json"),
			want:         contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"units":[1,2]}`), "sha256:source", "file:///source.json"),
			mutateOutput: true,
		},
		{
			name:        "mismatched content",
			sourceInput: contracts.NewLLMInputMaterial("source", "application/json", []byte(`{"units":[9]}`), "sha256:other", "file:///other.json"),
			wantErr:     "source input must match chunk",
		},
		{
			name: "default fields",
			sourceInput: contracts.LLMInputMaterial{
				Content:   []byte(`{"units":[1,2]}`),
				Digest:    "sha256:source",
				OriginURI: "file:///source.json",
			},
			want: contracts.LLMInputMaterial{
				Name:      "source",
				MediaType: "application/json",
				Content:   []byte(`{"units":[1,2]}`),
				Digest:    "sha256:source",
				OriginURI: "file:///source.json",
				SizeBytes: int64(len(`{"units":[1,2]}`)),
			},
		},
		{
			name: "preserve explicit metadata",
			sourceInput: contracts.LLMInputMaterial{
				Name:      "transcript",
				MediaType: "text/plain",
				Content:   []byte(`{"units":[1,2]}`),
				Digest:    "sha256:explicit",
				OriginURI: "file:///explicit.txt",
				SizeBytes: 42,
			},
			want: contracts.LLMInputMaterial{
				Name:      "transcript",
				MediaType: "text/plain",
				Content:   []byte(`{"units":[1,2]}`),
				Digest:    "sha256:explicit",
				OriginURI: "file:///explicit.txt",
				SizeBytes: 42,
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			req := extractionRequest()
			req.SourceInput = test.sourceInput
			if test.configure != nil {
				test.configure(&req)
			}
			ctx := test.ctx
			if ctx == nil && test.wantErr != "context must not be nil" {
				ctx = context.Background()
			}
			got, err := PrepareChunkExtraction(ctx, req)
			if test.wantErr != "" {
				if err == nil || !strings.Contains(err.Error(), test.wantErr) {
					t.Fatalf("PrepareChunkExtraction() error = %v, want %q", err, test.wantErr)
				}
				return
			}
			if err != nil {
				t.Fatalf("PrepareChunkExtraction() error = %v, want nil", err)
			}
			if !reflect.DeepEqual(got, test.want) {
				t.Fatalf("PrepareChunkExtraction() = %#v, want %#v", got, test.want)
			}
			if test.mutateOutput {
				got.Content[0] = 'x'
				if string(req.SourceInput.Content) != string(test.sourceInput.Content) || string(req.Chunk.Content) != `{"units":[1,2]}` {
					t.Fatal("PrepareChunkExtraction() output shares request content")
				}
			}
		})
	}
}

func extractionRequest() contracts.TypedExtractionRequest {
	doc := &source.SourceDocument{ID: "session-alpha"}
	chunk := &source.Chunk{
		ID:        "session-alpha:chunk:0",
		SourceID:  doc.ID,
		Content:   []byte(`{"units":[1,2]}`),
		MediaType: "application/json",
		Units:     []source.SourceUnit{{ID: 1}},
	}
	return contracts.TypedExtractionRequest{Source: doc, Chunk: chunk}
}