notarius/internal/modules/dnd/extract/locations/extractor_test.go

package locations

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

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity"
)

func TestExtractMapsLocationsWithOwnedEvidenceAndDeterministicOrder(t *testing.T) {
	client := &fakeLocationsLLMClient{response: extractionResponse{Locations: []locationResponse{
		{Name: "The Tavern", SourceRefs: responseSourceRefs(3, 3)},
		{Name: "Old Mill", SourceRefs: []locationSourceRefResponse{{StartUnitID: 2, EndUnitID: 2}, {StartUnitID: 1, EndUnitID: 1}, {StartUnitID: 1, EndUnitID: 1}}},
	}}}
	result, err := newExtractor(t, client).Extract(context.Background(), extractionRequest())
	if err != nil {
		t.Fatalf("Extract() error = %v, want nil", err)
	}
	refs := []source.SourceRef{{SourceID: "session-locations", StartUnitID: 1, EndUnitID: 1}, {SourceID: "session-locations", StartUnitID: 2, EndUnitID: 2}}
	want := dnd.LocationList{Locations: []dnd.Location{
		{ID: identity.DeriveID("Old Mill", refs), Name: "Old Mill", SourceRefs: refs},
		{ID: identity.DeriveID("The Tavern", []source.SourceRef{{SourceID: "session-locations", StartUnitID: 3, EndUnitID: 3}}), Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: "session-locations", StartUnitID: 3, EndUnitID: 3}}},
	}}
	if !reflect.DeepEqual(result.Value, want) {
		t.Fatalf("Value = %#v, want %#v", result.Value, want)
	}
	result.Value.Locations[0].SourceRefs[0].StartUnitID = 99
	for _, location := range client.response.Locations {
		for _, ref := range location.SourceRefs {
			if ref.StartUnitID == 99 {
				t.Fatal("result source references alias the model response")
			}
		}
	}
}

func TestExtractRetainsSameNameLocationsAtDifferentAnchors(t *testing.T) {
	client := &fakeLocationsLLMClient{response: extractionResponse{Locations: []locationResponse{
		{Name: "the tavern", SourceRefs: responseSourceRefs(1, 1)},
		{Name: "the tavern", SourceRefs: responseSourceRefs(3, 3)},
	}}}
	result, err := newExtractor(t, client).Extract(context.Background(), extractionRequest())
	if err != nil || len(result.Value.Locations) != 2 {
		t.Fatalf("Extract() = %#v, %v; want both same-name candidates", result, err)
	}
	if result.Value.Locations[0].ID == result.Value.Locations[1].ID || result.Value.Locations[0].Name != result.Value.Locations[1].Name {
		t.Fatalf("locations = %#v, want distinct evidence-anchored IDs", result.Value.Locations)
	}
}

func TestExtractPreservesInvalidCandidatesForValidators(t *testing.T) {
	client := &fakeLocationsLLMClient{content: []byte(`{"locations":[{"name":"","source_refs":[{"start_unit_id":0,"end_unit_id":-1}]}]}`)}
	result, err := newExtractor(t, client).Extract(context.Background(), extractionRequest())
	if err != nil {
		t.Fatalf("Extract() error = %v, want nil", err)
	}
	location := result.Value.Locations[0]
	if location.ID != "" || location.Name != "" || !reflect.DeepEqual(location.SourceRefs, []source.SourceRef{{SourceID: "session-locations", StartUnitID: 0, EndUnitID: -1}}) {
		t.Fatalf("location = %#v, want invalid candidate preserved", location)
	}
}

func TestExtractPassesReferencesWithoutTreatingThemAsEvidence(t *testing.T) {
	client := &fakeLocationsLLMClient{response: extractionResponse{Locations: []locationResponse{}}}
	req := extractionRequest()
	req.References = contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"glossary": {Slot: contracts.ReferenceSlot{Name: "glossary"}, Items: []contracts.ReferenceItem{{SlotName: "glossary", Content: []byte("Old Mill: abandoned granary")}}},
	}}
	if _, err := newExtractor(t, client).Extract(context.Background(), req); err != nil {
		t.Fatalf("Extract() error = %v", err)
	}
	inputs := client.requests[0].Inputs
	if string(inputs["glossary"].Content) != "Old Mill: abandoned granary" || strings.Contains(string(inputs["transcript"].Content), "abandoned granary") {
		t.Fatalf("prompt inputs = %#v, want separated reference material", inputs)
	}
}

func TestExtractDoesNotMutateRequestMaterials(t *testing.T) {
	client := &fakeLocationsLLMClient{response: extractionResponse{Locations: []locationResponse{{Name: "Old Mill", SourceRefs: responseSourceRefs(1, 1)}}}}
	req := extractionRequest()
	beforeUnits := append([]source.SourceUnit(nil), req.Source.Units...)
	beforeChunkUnits := append([]source.SourceUnit(nil), req.Chunk.Units...)
	beforeContent := append([]byte(nil), req.Chunk.Content...)
	if _, err := newExtractor(t, client).Extract(context.Background(), req); err != nil {
		t.Fatalf("Extract() error = %v", err)
	}
	if !reflect.DeepEqual(req.Source.Units, beforeUnits) || !reflect.DeepEqual(req.Chunk.Units, beforeChunkUnits) || !reflect.DeepEqual(req.Chunk.Content, beforeContent) {
		t.Fatalf("Extract() mutated request: %#v", req)
	}
}

func TestExtractHandlesEmptyOutputAndLocalFailures(t *testing.T) {
	empty, err := newExtractor(t, &fakeLocationsLLMClient{response: extractionResponse{Locations: []locationResponse{}}}).Extract(context.Background(), extractionRequest())
	if err != nil || empty.Value.Locations == nil || len(empty.Value.Locations) != 0 {
		t.Fatalf("empty Extract() = %#v, %v; want empty list", empty, err)
	}
	request := extractionRequest()
	var nilExtractor *Extractor
	for _, test := range []struct {
		name      string
		extractor *Extractor
		req       contracts.TypedExtractionRequest
		want      string
	}{
		{name: "nil extractor", extractor: nilExtractor, req: request, want: "extractor"},
		{name: "nil client", extractor: &Extractor{}, req: request, want: "LLM client"},
		{name: "preflight", extractor: newExtractor(t, &fakeLocationsLLMClient{}), req: mismatchedSourceInputRequest(request), want: "must match chunk"},
	} {
		t.Run(test.name, func(t *testing.T) {
			if _, err := test.extractor.Extract(context.Background(), test.req); err == nil || !strings.Contains(err.Error(), "dnd locations") || !strings.Contains(err.Error(), test.want) {
				t.Fatalf("Extract() error = %v, want local context", err)
			}
		})
	}
	_, err = newExtractor(t, &fakeLocationsLLMClient{err: errors.New("provider unavailable")}).Extract(context.Background(), request)
	if err == nil || !strings.Contains(err.Error(), "dnd locations") || !strings.Contains(err.Error(), "provider unavailable") {
		t.Fatalf("provider error = %v, want contextual provider error", err)
	}
}