312 lines
8.3 KiB
Go
312 lines
8.3 KiB
Go
package validate
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"io/fs"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"testing"
|
|
"testing/fstest"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
"github.com/santhosh-tekuri/jsonschema/v6"
|
|
)
|
|
|
|
func TestValidationCancellationBeforeWorkDoesNotOpenSchemaSource(t *testing.T) {
|
|
source := &countingSchemaFS{FS: fstest.MapFS{
|
|
"schema.json": {Data: []byte(`{"type":"object"}`)},
|
|
}}
|
|
validator := NewFSValidator(source, ".").(ValidationPreparer)
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
|
|
ValidationMode: domain.ValidationJSONSchema,
|
|
SchemaPath: "schema.json",
|
|
})
|
|
if plan != nil || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
|
|
}
|
|
if opens := source.opens.Load(); opens != 0 {
|
|
t.Fatalf("schema source opens = %d, want 0", opens)
|
|
}
|
|
}
|
|
|
|
func TestValidationCancellationBetweenReferencedSchemaReadChunks(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
source := &controlledSchemaFS{
|
|
FS: fstest.MapFS{
|
|
"root.json": {Data: []byte(`{"$ref":"child.json"}`)},
|
|
"child.json": {Data: []byte(`{"type":"string"}` + strings.Repeat(" ", schemaReadChunkSize*2))},
|
|
},
|
|
target: "child.json",
|
|
cancel: cancel,
|
|
}
|
|
validator := NewFSValidator(source, ".").(ValidationPreparer)
|
|
|
|
plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
|
|
ValidationMode: domain.ValidationJSONSchema,
|
|
SchemaPath: "root.json",
|
|
})
|
|
if plan != nil || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("PrepareValidation() = (%v, %v), want nil plan and context cancellation", plan, err)
|
|
}
|
|
if reads := source.reads.Load(); reads != 1 {
|
|
t.Fatalf("controlled child reads = %d, want 1", reads)
|
|
}
|
|
if closes := source.closes.Load(); closes != 1 {
|
|
t.Fatalf("controlled child closes = %d, want 1", closes)
|
|
}
|
|
}
|
|
|
|
func TestDecodeCancellationAfterSynchronousCallWins(t *testing.T) {
|
|
ctx := newCheckpointContext(2)
|
|
value, err := decodeJSONValue(ctx, []byte(`{"value":1}`))
|
|
if value != nil || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("decodeJSONValue() = (%v, %v), want nil value and context cancellation", value, err)
|
|
}
|
|
}
|
|
|
|
func TestCompileCancellationAfterSynchronousCallWins(t *testing.T) {
|
|
for _, dependencyErr := range []error{nil, errors.New("compile failed")} {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
compiled := &jsonschema.Schema{}
|
|
schema, err := compileJSONSchema(ctx, "promptkit-schema:/root.json", func(string) (*jsonschema.Schema, error) {
|
|
cancel()
|
|
return compiled, dependencyErr
|
|
})
|
|
if schema != nil || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("compileJSONSchema() = (%v, %v), want authoritative context cancellation", schema, err)
|
|
}
|
|
if dependencyErr != nil && errors.Is(err, dependencyErr) {
|
|
t.Fatalf("compileJSONSchema() error = %v, dependency error should not win", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestExecutionCancellationAfterSynchronousCallWins(t *testing.T) {
|
|
for _, dependencyErr := range []error{nil, errors.New("schema mismatch")} {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
executor := &controlledSchemaExecutor{cancel: cancel, err: dependencyErr}
|
|
validationErrors, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
|
|
if validationErrors != nil || !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("executeJSONSchema() = (%v, %v), want no result and context cancellation", validationErrors, err)
|
|
}
|
|
if calls := executor.calls.Load(); calls != 1 {
|
|
t.Fatalf("schema execution calls = %d, want 1", calls)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestCancellationDoesNotDetachBlockedSchemaRead(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
started := make(chan struct{})
|
|
release := make(chan struct{})
|
|
source := &controlledSchemaFS{
|
|
FS: fstest.MapFS{
|
|
"schema.json": {Data: []byte(`{"type":"object"}`)},
|
|
},
|
|
target: "schema.json",
|
|
started: started,
|
|
release: release,
|
|
}
|
|
validator := NewFSValidator(source, ".").(ValidationPreparer)
|
|
type outcome struct {
|
|
plan PreparedValidation
|
|
err error
|
|
}
|
|
result := make(chan outcome, 1)
|
|
go func() {
|
|
plan, err := validator.PrepareValidation(ctx, domain.OutputContract{
|
|
ValidationMode: domain.ValidationJSONSchema,
|
|
SchemaPath: "schema.json",
|
|
})
|
|
result <- outcome{plan: plan, err: err}
|
|
}()
|
|
|
|
waitForSignal(t, started, "schema read to start")
|
|
cancel()
|
|
select {
|
|
case got := <-result:
|
|
t.Fatalf("blocked dependency returned before release: (%v, %v)", got.plan, got.err)
|
|
default:
|
|
}
|
|
close(release)
|
|
got := waitForOutcome(t, result)
|
|
if got.plan != nil || !errors.Is(got.err, context.Canceled) {
|
|
t.Fatalf("PrepareValidation() after release = (%v, %v), want nil plan and context cancellation", got.plan, got.err)
|
|
}
|
|
if reads := source.reads.Load(); reads != 1 {
|
|
t.Fatalf("blocked reads = %d, want 1 completed read", reads)
|
|
}
|
|
}
|
|
|
|
func TestCancellationDoesNotDetachBlockedSchemaExecution(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
started := make(chan struct{})
|
|
release := make(chan struct{})
|
|
executor := &controlledSchemaExecutor{started: started, release: release}
|
|
result := make(chan error, 1)
|
|
go func() {
|
|
_, err := executeJSONSchema(ctx, executor, map[string]any{"ok": true})
|
|
result <- err
|
|
}()
|
|
|
|
waitForSignal(t, started, "schema execution to start")
|
|
cancel()
|
|
select {
|
|
case err := <-result:
|
|
t.Fatalf("blocked dependency returned before release: %v", err)
|
|
default:
|
|
}
|
|
close(release)
|
|
select {
|
|
case err := <-result:
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("executeJSONSchema() after release = %v, want context cancellation", err)
|
|
}
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timed out waiting for released schema execution")
|
|
}
|
|
if calls := executor.calls.Load(); calls != 1 {
|
|
t.Fatalf("schema execution calls = %d, want 1 completed call", calls)
|
|
}
|
|
}
|
|
|
|
type countingSchemaFS struct {
|
|
fs.FS
|
|
opens atomic.Int32
|
|
}
|
|
|
|
func (f *countingSchemaFS) Open(name string) (fs.File, error) {
|
|
f.opens.Add(1)
|
|
return f.FS.Open(name)
|
|
}
|
|
|
|
type controlledSchemaFS struct {
|
|
fs.FS
|
|
target string
|
|
cancel context.CancelFunc
|
|
started chan struct{}
|
|
release chan struct{}
|
|
once sync.Once
|
|
reads atomic.Int32
|
|
closes atomic.Int32
|
|
}
|
|
|
|
func (f *controlledSchemaFS) Open(name string) (fs.File, error) {
|
|
file, err := f.FS.Open(name)
|
|
if err != nil || name != f.target {
|
|
return file, err
|
|
}
|
|
return &controlledSchemaFile{File: file, owner: f}, nil
|
|
}
|
|
|
|
type controlledSchemaFile struct {
|
|
fs.File
|
|
owner *controlledSchemaFS
|
|
}
|
|
|
|
func (f *controlledSchemaFile) Read(buffer []byte) (int, error) {
|
|
f.owner.once.Do(func() {
|
|
if f.owner.started != nil {
|
|
close(f.owner.started)
|
|
}
|
|
if f.owner.release != nil {
|
|
<-f.owner.release
|
|
}
|
|
})
|
|
n, err := f.File.Read(buffer)
|
|
f.owner.reads.Add(1)
|
|
if f.owner.cancel != nil {
|
|
f.owner.cancel()
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func (f *controlledSchemaFile) Close() error {
|
|
f.owner.closes.Add(1)
|
|
return f.File.Close()
|
|
}
|
|
|
|
type controlledSchemaExecutor struct {
|
|
cancel context.CancelFunc
|
|
err error
|
|
started chan struct{}
|
|
release chan struct{}
|
|
calls atomic.Int32
|
|
}
|
|
|
|
func (e *controlledSchemaExecutor) Validate(any) error {
|
|
e.calls.Add(1)
|
|
if e.started != nil {
|
|
close(e.started)
|
|
}
|
|
if e.release != nil {
|
|
<-e.release
|
|
}
|
|
if e.cancel != nil {
|
|
e.cancel()
|
|
}
|
|
return e.err
|
|
}
|
|
|
|
type checkpointContext struct {
|
|
context.Context
|
|
mu sync.Mutex
|
|
remaining int
|
|
canceled bool
|
|
done chan struct{}
|
|
}
|
|
|
|
func newCheckpointContext(checksUntilCancel int) *checkpointContext {
|
|
return &checkpointContext{
|
|
Context: context.Background(),
|
|
remaining: checksUntilCancel,
|
|
done: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
func (c *checkpointContext) Done() <-chan struct{} {
|
|
return c.done
|
|
}
|
|
|
|
func (c *checkpointContext) Err() error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if c.canceled {
|
|
return context.Canceled
|
|
}
|
|
c.remaining--
|
|
if c.remaining == 0 {
|
|
c.canceled = true
|
|
close(c.done)
|
|
return context.Canceled
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func waitForSignal(t *testing.T, signal <-chan struct{}, description string) {
|
|
t.Helper()
|
|
select {
|
|
case <-signal:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatalf("timed out waiting for %s", description)
|
|
}
|
|
}
|
|
|
|
func waitForOutcome[T any](t *testing.T, result <-chan T) T {
|
|
t.Helper()
|
|
select {
|
|
case value := <-result:
|
|
return value
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("timed out waiting for released dependency")
|
|
var zero T
|
|
return zero
|
|
}
|
|
}
|