404 lines
11 KiB
Go
404 lines
11 KiB
Go
package contracts
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
)
|
|
|
|
var _ InputAdapter = fakeAdapter{}
|
|
var _ Chunker = fakeChunker{}
|
|
var _ Extractor = fakeExtractor{}
|
|
var _ Merger = fakeMerger{}
|
|
var _ Normalizer = fakeNormalizer{}
|
|
var _ Validator = fakeValidator{}
|
|
var _ StructuredLLMClient = fakeLLMClient{}
|
|
var _ OutputEncoder = fakeOutputEncoder{}
|
|
|
|
func TestFakeExtractorReturnsCandidateAndValidator(t *testing.T) {
|
|
validator := fakeValidator{name: "generic-validator"}
|
|
extractor := fakeExtractor{
|
|
key: "generic-extractor",
|
|
artifactType: "generic-artifact",
|
|
schemaVersion: "v1",
|
|
validators: []Validator{validator},
|
|
}
|
|
doc := &source.SourceDocument{
|
|
ID: "source-1",
|
|
Kind: "document",
|
|
Format: "text/plain",
|
|
Digest: "sha256:abc123",
|
|
Units: []source.SourceUnit{
|
|
{ID: "u1", Kind: "section", Text: "Source text."},
|
|
},
|
|
}
|
|
|
|
result, err := extractor.Extract(context.Background(), ExtractionRequest{Source: doc})
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
|
|
if extractor.Key() != "generic-extractor" {
|
|
t.Fatalf("Key() = %q, want generic-extractor", extractor.Key())
|
|
}
|
|
if extractor.ArtifactType() != "generic-artifact" {
|
|
t.Fatalf("ArtifactType() = %q, want generic-artifact", extractor.ArtifactType())
|
|
}
|
|
if extractor.SchemaVersion() != "v1" {
|
|
t.Fatalf("SchemaVersion() = %q, want v1", extractor.SchemaVersion())
|
|
}
|
|
if len(extractor.Validators()) != 1 {
|
|
t.Fatalf("len(Validators()) = %d, want 1", len(extractor.Validators()))
|
|
}
|
|
if extractor.Validators()[0].Name() != "generic-validator" {
|
|
t.Fatalf("Validators()[0].Name() = %q, want generic-validator", extractor.Validators()[0].Name())
|
|
}
|
|
if len(result.Candidates) != 1 {
|
|
t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates))
|
|
}
|
|
|
|
candidate := result.Candidates[0]
|
|
if candidate.Index != 0 {
|
|
t.Fatalf("ArtifactCandidate.Index = %d, want 0", candidate.Index)
|
|
}
|
|
if candidate.ExtractorKey != extractor.Key() {
|
|
t.Fatalf("ArtifactCandidate.ExtractorKey = %q, want %q", candidate.ExtractorKey, extractor.Key())
|
|
}
|
|
if candidate.ArtifactType != extractor.ArtifactType() {
|
|
t.Fatalf("ArtifactCandidate.ArtifactType = %q, want %q", candidate.ArtifactType, extractor.ArtifactType())
|
|
}
|
|
if candidate.SchemaVersion != extractor.SchemaVersion() {
|
|
t.Fatalf("ArtifactCandidate.SchemaVersion = %q, want %q", candidate.SchemaVersion, extractor.SchemaVersion())
|
|
}
|
|
if string(candidate.Payload) != `{"value":"example"}` {
|
|
t.Fatalf("ArtifactCandidate.Payload = %s, want example payload", candidate.Payload)
|
|
}
|
|
}
|
|
|
|
func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
|
|
doc := &source.SourceDocument{
|
|
ID: "source-1",
|
|
Kind: "document",
|
|
Format: "text/plain",
|
|
Digest: "sha256:abc123",
|
|
Units: []source.SourceUnit{
|
|
{ID: "u1", Kind: "section", Text: "Source text."},
|
|
},
|
|
}
|
|
chunker := fakeChunker{key: "generic-chunker"}
|
|
|
|
result, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc})
|
|
if err != nil {
|
|
t.Fatalf("Chunk() error = %v, want nil", err)
|
|
}
|
|
|
|
if chunker.Key() != "generic-chunker" {
|
|
t.Fatalf("Key() = %q, want generic-chunker", chunker.Key())
|
|
}
|
|
if len(result.Chunks) != 1 {
|
|
t.Fatalf("len(Chunks) = %d, want 1", len(result.Chunks))
|
|
}
|
|
|
|
chunk := result.Chunks[0]
|
|
if chunk.ID != "source-1:chunk:0" {
|
|
t.Fatalf("SourceChunk.ID = %q, want source-1:chunk:0", chunk.ID)
|
|
}
|
|
if chunk.SourceID != doc.ID {
|
|
t.Fatalf("SourceChunk.SourceID = %q, want %q", chunk.SourceID, doc.ID)
|
|
}
|
|
if chunk.Index != 0 {
|
|
t.Fatalf("SourceChunk.Index = %d, want 0", chunk.Index)
|
|
}
|
|
if len(chunk.Units) != 1 {
|
|
t.Fatalf("len(SourceChunk.Units) = %d, want 1", len(chunk.Units))
|
|
}
|
|
}
|
|
|
|
func TestFakeExtractorReceivesChunkAndAmbientContext(t *testing.T) {
|
|
extractor := fakeExtractor{
|
|
key: "generic-extractor",
|
|
artifactType: "generic-artifact",
|
|
schemaVersion: "v1",
|
|
}
|
|
doc := &source.SourceDocument{
|
|
ID: "source-1",
|
|
Kind: "document",
|
|
Format: "text/plain",
|
|
Digest: "sha256:abc123",
|
|
Units: []source.SourceUnit{
|
|
{ID: "u1", Kind: "section", Text: "First source text."},
|
|
{ID: "u2", Kind: "section", Text: "Second source text."},
|
|
},
|
|
}
|
|
chunk := SourceChunk{
|
|
ID: "source-1:chunk:1",
|
|
SourceID: doc.ID,
|
|
Index: 1,
|
|
Units: []source.SourceUnit{doc.Units[1]},
|
|
}
|
|
|
|
result, err := extractor.Extract(context.Background(), ExtractionRequest{
|
|
Source: doc,
|
|
Chunk: &chunk,
|
|
AmbientContext: map[string]any{"mode": "chunked"},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Extract() error = %v, want nil", err)
|
|
}
|
|
if len(result.Candidates) != 1 {
|
|
t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates))
|
|
}
|
|
|
|
candidate := result.Candidates[0]
|
|
if string(candidate.Payload) != `{"value":"chunked"}` {
|
|
t.Fatalf("ArtifactCandidate.Payload = %s, want chunked payload", candidate.Payload)
|
|
}
|
|
if len(candidate.SourceRefs) != 1 {
|
|
t.Fatalf("len(SourceRefs) = %d, want 1", len(candidate.SourceRefs))
|
|
}
|
|
ref := candidate.SourceRefs[0]
|
|
if ref.StartUnitID != "u2" || ref.EndUnitID != "u2" {
|
|
t.Fatalf("SourceRef = %+v, want u2 range", ref)
|
|
}
|
|
}
|
|
|
|
func TestFakeMergeNormalizeAndOutputContracts(t *testing.T) {
|
|
candidate := artifacts.ArtifactCandidate{
|
|
Index: 0,
|
|
ExtractorKey: "generic-extractor",
|
|
ArtifactType: "generic-artifact",
|
|
SchemaVersion: "v1",
|
|
Payload: json.RawMessage(`{"value":"example"}`),
|
|
}
|
|
chunk := SourceChunk{
|
|
ID: "source-1:chunk:0",
|
|
SourceID: "source-1",
|
|
Index: 0,
|
|
Units: []source.SourceUnit{
|
|
{ID: "u1", Kind: "section", Text: "Source text."},
|
|
},
|
|
}
|
|
merger := fakeMerger{key: "generic-merger"}
|
|
normalizer := fakeNormalizer{key: "generic-normalizer"}
|
|
encoder := fakeOutputEncoder{key: "generic-output"}
|
|
|
|
merged, err := merger.Merge(context.Background(), MergeRequest{
|
|
LaneID: "generic-artifact",
|
|
ChunkArtifacts: []ChunkArtifacts{
|
|
{
|
|
Chunk: chunk,
|
|
Candidates: []artifacts.ArtifactCandidate{candidate},
|
|
},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Merge() error = %v, want nil", err)
|
|
}
|
|
if merger.Key() != "generic-merger" {
|
|
t.Fatalf("Merger.Key() = %q, want generic-merger", merger.Key())
|
|
}
|
|
if len(merged.Candidates) != 1 {
|
|
t.Fatalf("len(merged.Candidates) = %d, want 1", len(merged.Candidates))
|
|
}
|
|
|
|
normalized, err := normalizer.Normalize(context.Background(), NormalizeRequest{
|
|
LaneID: "generic-artifact",
|
|
Candidates: merged.Candidates,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Normalize() error = %v, want nil", err)
|
|
}
|
|
if normalizer.Key() != "generic-normalizer" {
|
|
t.Fatalf("Normalizer.Key() = %q, want generic-normalizer", normalizer.Key())
|
|
}
|
|
if len(normalized.Candidates) != 1 {
|
|
t.Fatalf("len(normalized.Candidates) = %d, want 1", len(normalized.Candidates))
|
|
}
|
|
|
|
encoded, err := encoder.Encode(context.Background(), OutputRequest{
|
|
Manifest: artifacts.RunManifest{RunID: "run-1"},
|
|
Approved: []artifacts.Artifact{
|
|
artifacts.ArtifactFromCandidate(normalized.Candidates[0]),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Encode() error = %v, want nil", err)
|
|
}
|
|
if encoder.Key() != "generic-output" {
|
|
t.Fatalf("OutputEncoder.Key() = %q, want generic-output", encoder.Key())
|
|
}
|
|
if encoded.ContentType != "application/json" {
|
|
t.Fatalf("ContentType = %q, want application/json", encoded.ContentType)
|
|
}
|
|
if string(encoded.Bytes) != `{"run_id":"run-1","approved_count":1}` {
|
|
t.Fatalf("Bytes = %s, want encoded output", encoded.Bytes)
|
|
}
|
|
}
|
|
|
|
type fakeAdapter struct {
|
|
key string
|
|
doc *source.SourceDocument
|
|
}
|
|
|
|
func (adapter fakeAdapter) Key() string {
|
|
return adapter.key
|
|
}
|
|
|
|
func (adapter fakeAdapter) Parse(ctx context.Context, req ParseRequest) (*source.SourceDocument, error) {
|
|
return adapter.doc, nil
|
|
}
|
|
|
|
type fakeChunker struct {
|
|
key string
|
|
}
|
|
|
|
func (chunker fakeChunker) Key() string {
|
|
return chunker.key
|
|
}
|
|
|
|
func (chunker fakeChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
|
|
return ChunkResult{
|
|
Chunks: []SourceChunk{
|
|
{
|
|
ID: req.Source.ID + ":chunk:0",
|
|
SourceID: req.Source.ID,
|
|
Index: 0,
|
|
Units: append([]source.SourceUnit(nil), req.Source.Units...),
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type fakeExtractor struct {
|
|
key string
|
|
artifactType string
|
|
schemaVersion string
|
|
validators []Validator
|
|
}
|
|
|
|
func (extractor fakeExtractor) Key() string {
|
|
return extractor.key
|
|
}
|
|
|
|
func (extractor fakeExtractor) ArtifactType() string {
|
|
return extractor.artifactType
|
|
}
|
|
|
|
func (extractor fakeExtractor) SchemaVersion() string {
|
|
return extractor.schemaVersion
|
|
}
|
|
|
|
func (extractor fakeExtractor) Validators() []Validator {
|
|
return extractor.validators
|
|
}
|
|
|
|
func (extractor fakeExtractor) Extract(ctx context.Context, req ExtractionRequest) (ExtractionResult, error) {
|
|
units := req.Source.Units
|
|
if req.Chunk != nil {
|
|
units = req.Chunk.Units
|
|
}
|
|
payload := json.RawMessage(`{"value":"example"}`)
|
|
if req.AmbientContext["mode"] == "chunked" {
|
|
payload = json.RawMessage(`{"value":"chunked"}`)
|
|
}
|
|
|
|
return ExtractionResult{
|
|
Candidates: []artifacts.ArtifactCandidate{
|
|
{
|
|
Index: 0,
|
|
ExtractorKey: extractor.key,
|
|
ArtifactType: extractor.artifactType,
|
|
SchemaVersion: extractor.schemaVersion,
|
|
Payload: payload,
|
|
SourceRefs: []source.SourceRef{
|
|
{
|
|
SourceID: req.Source.ID,
|
|
StartUnitID: units[0].ID,
|
|
EndUnitID: units[len(units)-1].ID,
|
|
},
|
|
},
|
|
},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type fakeMerger struct {
|
|
key string
|
|
}
|
|
|
|
func (merger fakeMerger) Key() string {
|
|
return merger.key
|
|
}
|
|
|
|
func (merger fakeMerger) Merge(ctx context.Context, req MergeRequest) (MergeResult, error) {
|
|
var candidates []artifacts.ArtifactCandidate
|
|
for _, chunkArtifacts := range req.ChunkArtifacts {
|
|
candidates = append(candidates, chunkArtifacts.Candidates...)
|
|
}
|
|
|
|
return MergeResult{Candidates: candidates}, nil
|
|
}
|
|
|
|
type fakeNormalizer struct {
|
|
key string
|
|
}
|
|
|
|
func (normalizer fakeNormalizer) Key() string {
|
|
return normalizer.key
|
|
}
|
|
|
|
func (normalizer fakeNormalizer) Normalize(ctx context.Context, req NormalizeRequest) (NormalizeResult, error) {
|
|
return NormalizeResult{Candidates: req.Candidates}, nil
|
|
}
|
|
|
|
type fakeValidator struct {
|
|
name string
|
|
}
|
|
|
|
func (validator fakeValidator) Name() string {
|
|
return validator.name
|
|
}
|
|
|
|
func (validator fakeValidator) Validate(ctx context.Context, req ValidationRequest) (ValidationResult, error) {
|
|
decisions := make([]ValidationDecision, 0, len(req.Candidates))
|
|
for _, candidate := range req.Candidates {
|
|
decisions = append(decisions, ValidationDecision{
|
|
CandidateIndex: candidate.Index,
|
|
Approved: true,
|
|
ReasonCode: "accepted",
|
|
Message: "candidate accepted",
|
|
})
|
|
}
|
|
|
|
return ValidationResult{
|
|
ValidatorName: validator.name,
|
|
Decisions: decisions,
|
|
}, nil
|
|
}
|
|
|
|
type fakeLLMClient struct{}
|
|
|
|
func (client fakeLLMClient) CompleteStructured(ctx context.Context, req StructuredCompletionRequest, out any) (StructuredCompletionResponse, error) {
|
|
return StructuredCompletionResponse{
|
|
Content: json.RawMessage(`{"value":"example"}`),
|
|
}, nil
|
|
}
|
|
|
|
type fakeOutputEncoder struct {
|
|
key string
|
|
}
|
|
|
|
func (encoder fakeOutputEncoder) Key() string {
|
|
return encoder.key
|
|
}
|
|
|
|
func (encoder fakeOutputEncoder) Encode(ctx context.Context, req OutputRequest) (OutputResult, error) {
|
|
return OutputResult{
|
|
Bytes: []byte(`{"run_id":"` + req.Manifest.RunID + `","approved_count":1}`),
|
|
ContentType: "application/json",
|
|
}, nil
|
|
}
|