344 lines
9.9 KiB
Go
344 lines
9.9 KiB
Go
package contracts
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"reflect"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
)
|
|
|
|
var _ InputAdapter = fakeAdapter{}
|
|
var _ Chunker = fakeChunker{}
|
|
var _ Extractor[fakeArtifact] = fakeExtractor{}
|
|
var _ Merger[fakeArtifact] = fakeMerger{}
|
|
var _ Normalizer[fakeArtifact] = fakeNormalizer{}
|
|
var _ StructuredLLMClient = fakeLLMClient{}
|
|
var _ OutputEncoder = fakeOutputEncoder{}
|
|
|
|
func TestCloneReferenceSlotsEmptyInputReturnsNil(t *testing.T) {
|
|
if got := CloneReferenceSlots(nil); got != nil {
|
|
t.Fatalf("CloneReferenceSlots(nil) = %#v, want nil", got)
|
|
}
|
|
if got := CloneReferenceSlots([]ReferenceSlot{}); got != nil {
|
|
t.Fatalf("CloneReferenceSlots(empty) = %#v, want nil", got)
|
|
}
|
|
}
|
|
|
|
func TestCloneSerializedArtifactPointerOwnsArtifactData(t *testing.T) {
|
|
original := &SerializedArtifact{
|
|
Kind: "test/artifact",
|
|
Schema: ArtifactSchema{ID: "test.schema", JSONSchema: []byte(`{"type":"object"}`)},
|
|
Content: []byte(`{"items":[]}`),
|
|
Metadata: map[string]any{"count": 1},
|
|
}
|
|
cloned := CloneSerializedArtifactPointer(original)
|
|
original.Schema.JSONSchema[0] = '['
|
|
original.Content[0] = '['
|
|
original.Metadata["count"] = 2
|
|
if cloned == nil || string(cloned.Schema.JSONSchema) != `{"type":"object"}` || string(cloned.Content) != `{"items":[]}` || cloned.Metadata["count"] != 1 {
|
|
t.Fatalf("CloneSerializedArtifactPointer() = %#v, want independent artifact data", cloned)
|
|
}
|
|
if CloneSerializedArtifactPointer(nil) != nil {
|
|
t.Fatal("CloneSerializedArtifactPointer(nil) must return nil")
|
|
}
|
|
}
|
|
|
|
func TestCloneReferenceSlotsPreservesFields(t *testing.T) {
|
|
slots := []ReferenceSlot{
|
|
{
|
|
Name: "roster",
|
|
Description: "Known characters",
|
|
Required: true,
|
|
AcceptedMediaTypes: []string{"text/plain", "text/markdown"},
|
|
Multiple: true,
|
|
MaxBytes: 4096,
|
|
},
|
|
{
|
|
Name: "glossary",
|
|
Description: "Campaign terms",
|
|
MaxBytes: 2048,
|
|
},
|
|
}
|
|
|
|
got := CloneReferenceSlots(slots)
|
|
|
|
if !reflect.DeepEqual(got, slots) {
|
|
t.Fatalf("CloneReferenceSlots() = %#v, want %#v", got, slots)
|
|
}
|
|
}
|
|
|
|
func TestCloneReferenceSlotsCopiesAcceptedMediaTypes(t *testing.T) {
|
|
slots := []ReferenceSlot{
|
|
{
|
|
Name: "party",
|
|
AcceptedMediaTypes: []string{"application/json", "text/plain"},
|
|
},
|
|
}
|
|
|
|
got := CloneReferenceSlots(slots)
|
|
got[0].Name = "changed"
|
|
got[0].AcceptedMediaTypes[0] = "text/markdown"
|
|
|
|
if slots[0].Name != "party" {
|
|
t.Fatalf("source slot name = %q, want unchanged", slots[0].Name)
|
|
}
|
|
if slots[0].AcceptedMediaTypes[0] != "application/json" {
|
|
t.Fatalf("source AcceptedMediaTypes aliased clone: %#v", slots[0].AcceptedMediaTypes)
|
|
}
|
|
}
|
|
|
|
func TestCloneReferenceSlotsCopiesAcceptedArtifactKinds(t *testing.T) {
|
|
slots := []ReferenceSlot{{Name: "npcs", AcceptedArtifactKinds: []ArtifactKind{"dnd/npc-registry"}}}
|
|
clone := CloneReferenceSlots(slots)
|
|
clone[0].AcceptedArtifactKinds[0] = "changed"
|
|
if slots[0].AcceptedArtifactKinds[0] != "dnd/npc-registry" {
|
|
t.Fatalf("source AcceptedArtifactKinds aliased clone: %#v", slots[0].AcceptedArtifactKinds)
|
|
}
|
|
}
|
|
|
|
func TestReferenceItemJSONOmitsContent(t *testing.T) {
|
|
item := ReferenceItem{
|
|
SlotName: "roster",
|
|
MediaType: "text/plain",
|
|
Content: []byte("reference content"),
|
|
Digest: "sha256:reference",
|
|
Origin: ReferenceOrigin{Type: "file", URI: "file:///tmp/roster.txt"},
|
|
}
|
|
|
|
encoded, err := json.Marshal(item)
|
|
if err != nil {
|
|
t.Fatalf("json.Marshal() error = %v, want nil", err)
|
|
}
|
|
|
|
var got map[string]any
|
|
if err := json.Unmarshal(encoded, &got); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v, want nil", err)
|
|
}
|
|
if _, ok := got["content"]; ok {
|
|
t.Fatalf("encoded reference item leaked content: %s", encoded)
|
|
}
|
|
if _, ok := got["Content"]; ok {
|
|
t.Fatalf("encoded reference item leaked Content: %s", encoded)
|
|
}
|
|
}
|
|
|
|
func TestLLMInputMaterialCopiesContentAndOmitsContentFromJSON(t *testing.T) {
|
|
content := []byte("raw source bytes")
|
|
material := NewLLMInputMaterial("transcript", "application/json", content, "sha256:source", "file:///tmp/source.json")
|
|
content[0] = 'R'
|
|
if got := string(material.Content); got != "raw source bytes" {
|
|
t.Fatalf("material content = %q, want defensive copy", got)
|
|
}
|
|
if material.SizeBytes != int64(len("raw source bytes")) {
|
|
t.Fatalf("SizeBytes = %d, want content length", material.SizeBytes)
|
|
}
|
|
|
|
clone := material.Clone()
|
|
clone.Content[0] = 'X'
|
|
if got := string(material.Content); got != "raw source bytes" {
|
|
t.Fatalf("cloned material content aliased original: %q", got)
|
|
}
|
|
|
|
encoded, err := json.Marshal(material)
|
|
if err != nil {
|
|
t.Fatalf("json.Marshal() error = %v, want nil", err)
|
|
}
|
|
var got map[string]any
|
|
if err := json.Unmarshal(encoded, &got); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v, want nil", err)
|
|
}
|
|
if _, ok := got["content"]; ok {
|
|
t.Fatalf("encoded material leaked content: %s", encoded)
|
|
}
|
|
if _, ok := got["Content"]; ok {
|
|
t.Fatalf("encoded material leaked Content: %s", encoded)
|
|
}
|
|
if got["digest"] != "sha256:source" || got["origin_uri"] != "file:///tmp/source.json" {
|
|
t.Fatalf("encoded material = %#v, want non-secret provenance", got)
|
|
}
|
|
}
|
|
|
|
func TestLLMInputSetCloneCopiesContent(t *testing.T) {
|
|
set := LLMInputSet{
|
|
"transcript": NewLLMInputMaterial("transcript", "application/json", []byte("source"), "sha256:source", "file:///tmp/source.json"),
|
|
}
|
|
clone := set.Clone()
|
|
clone["transcript"].Content[0] = 'S'
|
|
if got := string(set["transcript"].Content); got != "source" {
|
|
t.Fatalf("input set clone aliased content: %q", got)
|
|
}
|
|
}
|
|
|
|
func TestArtifactSchemaJSONOmitsSchemaContent(t *testing.T) {
|
|
schema := ArtifactSchema{
|
|
ID: "schema-id",
|
|
Name: "schema-name",
|
|
Version: "v1",
|
|
JSONSchema: []byte(`{"type":"object"}`),
|
|
}
|
|
encoded, err := json.Marshal(schema)
|
|
if err != nil {
|
|
t.Fatalf("json.Marshal() error = %v, want nil", err)
|
|
}
|
|
var got map[string]any
|
|
if err := json.Unmarshal(encoded, &got); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v, want nil", err)
|
|
}
|
|
if got["id"] != "schema-id" || got["name"] != "schema-name" || got["version"] != "v1" {
|
|
t.Fatalf("encoded schema = %#v, want schema provenance", got)
|
|
}
|
|
if _, ok := got["json_schema"]; ok {
|
|
t.Fatalf("encoded schema leaked raw schema content: %s", encoded)
|
|
}
|
|
if _, ok := got["JSONSchema"]; ok {
|
|
t.Fatalf("encoded schema leaked raw schema content: %s", encoded)
|
|
}
|
|
}
|
|
|
|
func TestOutputFileJSONShapeOmitsBytes(t *testing.T) {
|
|
file := OutputFile{
|
|
Name: "artifacts/events.json",
|
|
ContentType: "application/json",
|
|
Bytes: []byte(`{"ignored":true}`),
|
|
}
|
|
|
|
encoded, err := json.Marshal(file)
|
|
if err != nil {
|
|
t.Fatalf("json.Marshal() error = %v, want nil", err)
|
|
}
|
|
|
|
var got map[string]any
|
|
if err := json.Unmarshal(encoded, &got); err != nil {
|
|
t.Fatalf("json.Unmarshal() error = %v, want nil", err)
|
|
}
|
|
if got["name"] != "artifacts/events.json" {
|
|
t.Fatalf("name = %#v, want logical file name", got["name"])
|
|
}
|
|
if got["content_type"] != "application/json" {
|
|
t.Fatalf("content_type = %#v, want application/json", got["content_type"])
|
|
}
|
|
if _, ok := got["Bytes"]; ok {
|
|
t.Fatalf("encoded output file leaked Bytes: %s", encoded)
|
|
}
|
|
if _, ok := got["bytes"]; ok {
|
|
t.Fatalf("encoded output file leaked bytes: %s", encoded)
|
|
}
|
|
}
|
|
|
|
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) ReferenceSlots() []ReferenceSlot {
|
|
return nil
|
|
}
|
|
|
|
func (chunker fakeChunker) Plan(ctx context.Context, req ChunkRequest) (ChunkPlanResult, error) {
|
|
return ChunkPlanResult{
|
|
Plan: source.ChunkPlan{
|
|
SourceDigest: req.Source.Digest,
|
|
Ranges: []source.ChunkRange{{
|
|
StartUnitID: req.Source.Units[0].ID,
|
|
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
|
}},
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type fakeExtractor struct {
|
|
key string
|
|
}
|
|
|
|
type fakeArtifact struct{ Value string }
|
|
|
|
func (extractor fakeExtractor) Key() string {
|
|
return extractor.key
|
|
}
|
|
|
|
func (extractor fakeExtractor) ReferenceSlots() []ReferenceSlot {
|
|
return nil
|
|
}
|
|
|
|
func (extractor fakeExtractor) Extract(ctx context.Context, req TypedExtractionRequest) (TypedExtractionResult[fakeArtifact], error) {
|
|
value := "example"
|
|
if req.AmbientContext["mode"] == "chunked" {
|
|
value = "chunked"
|
|
}
|
|
return TypedExtractionResult[fakeArtifact]{Value: fakeArtifact{Value: value}}, nil
|
|
}
|
|
|
|
type fakeMerger struct {
|
|
key string
|
|
}
|
|
|
|
func (merger fakeMerger) Key() string {
|
|
return merger.key
|
|
}
|
|
|
|
func (merger fakeMerger) Merge(ctx context.Context, req TypedMergeRequest[fakeArtifact]) (TypedMergeResult[fakeArtifact], error) {
|
|
return TypedMergeResult[fakeArtifact]{Value: req.ExtractOutputs[0].Value}, nil
|
|
}
|
|
|
|
type fakeNormalizer struct {
|
|
key string
|
|
}
|
|
|
|
func (normalizer fakeNormalizer) Key() string {
|
|
return normalizer.key
|
|
}
|
|
|
|
func (normalizer fakeNormalizer) ReferenceSlots() []ReferenceSlot {
|
|
return nil
|
|
}
|
|
|
|
func (normalizer fakeNormalizer) Normalize(ctx context.Context, req TypedNormalizeRequest[fakeArtifact]) (TypedNormalizeResult[fakeArtifact], error) {
|
|
return TypedNormalizeResult[fakeArtifact]{Value: req.MergeOutput.Value}, 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{
|
|
Files: []OutputFile{
|
|
{
|
|
Name: "artifacts/generic.json",
|
|
ContentType: "application/json",
|
|
Bytes: []byte(`{"run_id":"` + req.Manifest.RunID + `","output_count":1}`),
|
|
},
|
|
},
|
|
}, nil
|
|
}
|