Files
notarius/internal/modules/dnd/normalize/locationregistry/test_helpers_test.go

65 lines
2.3 KiB
Go

package locationregistry
import (
"context"
"encoding/json"
"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"
)
type recordingLocationNormalizerClient struct {
response string
err error
requests []contracts.StructuredCompletionRequest
}
func (c *recordingLocationNormalizerClient) CompleteStructured(_ context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) {
c.requests = append(c.requests, request)
if c.err != nil {
return contracts.StructuredCompletionResponse{}, c.err
}
response := c.response
if response == "" {
response = `{"duplicate_groups":[]}`
}
content := []byte(response)
if err := json.Unmarshal(content, output); err != nil {
return contracts.StructuredCompletionResponse{}, err
}
return contracts.StructuredCompletionResponse{Content: content}, nil
}
func newNormalizer(t *testing.T, client contracts.StructuredLLMClient) *Normalizer {
t.Helper()
normalizer, err := New(client, Options{})
if err != nil {
t.Fatalf("New() error = %v", err)
}
return normalizer
}
func normalizeRequest(value dnd.LocationRegistry) contracts.TypedNormalizeRequest[dnd.LocationRegistry] {
return contracts.TypedNormalizeRequest[dnd.LocationRegistry]{
Source: &source.SourceDocument{},
MergeOutput: contracts.MergeArtifact[dnd.LocationRegistry]{Value: value},
}
}
func normalizeRequestWithSource(value dnd.LocationRegistry, doc *source.SourceDocument) contracts.TypedNormalizeRequest[dnd.LocationRegistry] {
request := normalizeRequest(value)
request.Source = doc
return request
}
func semanticDocument() *source.SourceDocument {
return &source.SourceDocument{ID: "location-session", Units: []source.SourceUnit{{ID: 10, Kind: "speech", Text: "The old mill is the Greencloak's refuge."}, {ID: 20, Kind: "speech", Text: "The mill stands on the northern road."}, {ID: 30, Kind: "speech", Text: "The tavern is beside the mill."}, {ID: 40, Kind: "speech", Text: "The mill's cellar is flooded."}}}
}
func hasWarning(warnings []contracts.Warning, reason string) bool {
for _, warning := range warnings {
if warning.ReasonCode == reason {
return true
}
}
return false
}