Test continuous execution budgets across collaborators

internal/usecase/operation_budget_test.go

@@ -0,0 +1,54 @@
+package usecase
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"testing/synctest"
+	"time"
+
+	"go.maximumdirect.net/promptkit/internal/domain"
+	"go.maximumdirect.net/promptkit/internal/validate"
+)
+
+type deadlineValidation struct{ deadline time.Time }
+
+func (v *deadlineValidation) PrepareValidation(ctx context.Context, _ domain.OutputContract) (validate.PreparedValidation, error) {
+	v.deadline, _ = ctx.Deadline()
+	return v, nil
+}
+func (v *deadlineValidation) Validate(ctx context.Context, _ *domain.Artifact) (domain.ValidationResult, error) {
+	deadline, _ := ctx.Deadline()
+	if deadline != v.deadline {
+		return domain.ValidationResult{}, errors.New("validation deadline changed")
+	}
+	<-ctx.Done()
+	return domain.ValidationResult{}, ctx.Err()
+}
+func (*deadlineValidation) SchemaDocument() any { return nil }
+
+type deadlineModel struct{ validation *deadlineValidation }
+
+func (m deadlineModel) Generate(ctx context.Context, _ domain.GenerateRequest) (*domain.GenerateResponse, error) {
+	deadline, _ := ctx.Deadline()
+	if deadline != m.validation.deadline {
+		return nil, errors.New("generation deadline changed")
+	}
+	return &domain.GenerateResponse{Content: "ok"}, nil
+}
+
+func TestOperationDeadlineReachesPreparationGenerationAndValidation(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		validation := &deadlineValidation{}
+		runner := NewRunner(&fakePromptRepo{def: promptDef(domain.FormatText, domain.ValidationBasic, 0)},
+			&fakeExecutionProfileRepo{profiles: map[string]*domain.ExecutionProfile{"exec": defaultExecutionProfile()}},
+			defaultArtifactReader(), defaultRenderer(), deadlineModel{validation}, validation, nil)
+		scope := NewOperationScope(context.Background(), time.Second)
+		defer scope.Close()
+		_, err := runner.Run(scope.Context, domain.RunRequest{PromptID: "p", ProfileID: "exec", Inputs: singleInputRef()})
+		err = scope.Finalize(err)
+		if !errors.Is(err, ErrOperationTimeout) || !errors.Is(err, ErrValidation) || !errors.Is(err, context.DeadlineExceeded) {
+			t.Fatalf("validation error=%v", err)
+		}
+	})
+}