Add Override, to replace the task which provides a key

README.md

@@ -27,6 +27,49 @@ result, execution of the graph is performed as a parallelised breadth first sear
 This package contains some creative APIs which work around limitations in Golang generics; chief
 of these is the `ID` type used in place of slices of keys because Golang slices are invariant.
 
+## Overriding tasks in tests
+
+A graph's tasks often create things for themselves: clients of other services, say. That makes a
+whole graph hard to test, as there is no way to give it a fake client. `taskgraphtest` can replace
+the task which provides a key:
+
+```go
+taskgraphtest.RunWithOverrides(ctx, graph, inputs, []tg.Binding{keyClient.Bind(fakeClient)})
+```
+
+The task providing `keyClient` is not run, and the tasks which depend on it see the fake. This makes
+a graph's dependencies a matter of which keys its tasks provide: anything which a test may want to
+replace should be a key of its own.
+
+Replacing tasks is only available from `taskgraphtest`. `Graph.Run` has no way to do it, and the
+engine's support is in an `internal` package, which nothing outside this module can import. Code which
+runs a graph and needs to be tested can take a function for running graphs, and be handed
+`Overrider.Run` by its tests.
+
+* An input for a key which a task provides is still an error, so a key is only ever replaced on
+  purpose. The same goes for overriding one key twice, or overriding a key which is also an input.
+* If a task provides several keys they must all be overridden, or none (`ErrPartialOverride`).
+  Otherwise the task would still have to run, with all of its effects, with some of its results
+  thrown away.
+* Overriding a key which no task provides is an error (`ErrUnknownOverride`), which catches
+  overrides left behind by a task which has since been renamed or removed. `CheckWithOverrides`
+  finds these without running the graph.
+* The override may be absent: `keyClient.BindError(err)` replaces the task with a failure, to see how
+  the tasks which depend on it cope.
+* The overrides are included in the `Binder` returned, as the results of the tasks they replace would
+  have been.
+* Inputs which only a replaced task needed are not required.
+* Tasks which only produced inputs for a replaced task still run, as the graph can't know that its
+  caller doesn't want their results.
+* Overrides apply to the graph they are given. A graph run as a task with `AsTask` is replaced as a
+  whole, by overriding every key it exposes, but its tasks can't be overridden individually.
+* An `Overrider` holds one set of overrides for something which runs several graphs, such as an
+  operator with a graph to create and a graph to delete. Each graph is run with the overrides for the
+  keys it provides, and `Overrider.Check` makes sure that every override is for some graph.
+
+`taskgraphtest.Test` and `taskgraphtest.Suite` take `Overrides`, those of a `Test` replacing those
+of its `Suite` for the same key.
+
 ## Limitations and sharp edges
 
 * Taskgraph has no deadlock detection. In theory the fact that the graph is acyclic should prevent

graph.go

@@ -90,9 +90,16 @@ type Graph interface {
 type runState struct {
 	Binder
 	signals map[string]chan struct{}
+	// skipped holds the IDs of the nodes which are not run because their tasks are overridden. They
+	// have nothing listening for signals, and nothing which depends on them waits for them.
+	skipped map[string]struct{}
 }
 
 func (rs *runState) signal(ctx context.Context, childID string) (err error) {
+	if _, skipped := rs.skipped[childID]; skipped {
+		return nil
+	}
+
 	// Capture (and ignore) panics caused by sends to closed channels, which may occur if a task
 	// starts slowly and sees that all of its dependencies are available before we can signal it.
 	defer func() {
@@ -127,6 +134,8 @@ type graphNode struct {
 
 const (
 	traceTaskgraphAbsentKeysPrefix = "taskgraph.absent_keys."
+	traceTaskgraphOverriddenKeys   = "taskgraph.overridden_keys"
+	traceTaskgraphSkippedTasks     = "taskgraph.skipped_tasks"
 )
 
 // Execute the task against the binder provided in the runState.
@@ -262,34 +271,24 @@ type graph struct {
 	logger                       Logger
 }
 
-func (g *graph) buildInputBinder(inputs ...Binding) (Binder, error) {
-	b := NewBinder()
-
-	if err := b.Store(inputs...); err != nil {
-		return nil, wrapStackErrorf("duplicate input: %w", err)
-	}
-
-	var missingInputs []string
-	for requiredInput := range g.allDependencies.Difference(g.allProvided).Iter() {
-		if !b.Has(requiredInput) {
-			missingInputs = append(missingInputs, requiredInput.String())
-		}
-	}
-	if len(missingInputs) > 0 {
-		return nil, wrapStackErrorf("%w: %s", ErrMissingInputs, strings.Join(missingInputs, ", "))
-	}
-
-	return b, nil
-}
-
 // Check is Graph.Check.
 func (g *graph) Check(inputs ...Binding) error {
-	_, err := g.buildInputBinder(inputs...)
+	_, err := g.buildInputs(inputs, nil, false)
 	return err
 }
 
 // Run is Graph.Run.
-func (g *graph) Run(ctx context.Context, inputs ...Binding) (b Binder, err error) {
+func (g *graph) Run(ctx context.Context, inputs ...Binding) (Binder, error) {
+	return g.runWithOverrides(ctx, inputs, nil, false)
+}
+
+// runWithOverrides runs the graph with some of its tasks replaced by the overrides, which are for
+// tests (see taskgraphtest).
+func (g *graph) runWithOverrides(
+	ctx context.Context,
+	inputs, overrides []Binding,
+	ignoreUnprovided bool,
+) (b Binder, err error) {
 	startTime := time.Now()
 	defer func() {
 		result := "success"
@@ -299,20 +298,33 @@ func (g *graph) Run(ctx context.Context, inputs ...Binding) (b Binder, err error
 		executionLatency.WithLabelValues(g.name, result).
 			Observe(float64(time.Since(startTime) / time.Millisecond))
 	}()
-	base, err := g.buildInputBinder(inputs...)
+	in, err := g.buildInputs(inputs, overrides, ignoreUnprovided)
 	if err != nil {
 		return nil, err
 	}
 
+	// Overrides stand in for the results of the tasks they replace, so they are stored with the
+	// results, and are returned along with them.
 	outputs := NewBinder()
+	if err := outputs.Store(in.overrides...); err != nil {
+		return nil, wrapStackErrorf("storing overrides: %w", err)
+	}
 	overlay := &overlayBinder{
-		base:    base,
+		base:    in.inputs,
 		overlay: outputs,
 	}
 
 	tCtx, span := g.tracer.Start(ctx, g.name)
 	defer span.End()
-	if err := g.runWithBinder(tCtx, overlay); err != nil {
+	if len(in.overrides) > 0 {
+		span.SetAttributes(
+			attribute.StringSlice(traceTaskgraphOverriddenKeys, in.overriddenKeys()),
+			attribute.StringSlice(traceTaskgraphSkippedTasks, in.skippedTaskNames()),
+		)
+		g.logger.Debugf("graph %s overriding keys %v, so not running tasks %v",
+			g.name, in.overriddenKeys(), in.skippedTaskNames())
+	}
+	if err := g.runWithBinder(tCtx, overlay, in.skipped); err != nil {
 		span.RecordError(err)
 		return nil, err
 	}
@@ -322,20 +334,31 @@ func (g *graph) Run(ctx context.Context, inputs ...Binding) (b Binder, err error
 
 // Sets up the per-run state of the graph, and runs all of the tasks in their own goroutines until
 // all have terminated. If any task returns an error, the entire graph run is cancelled.
-func (g *graph) runWithBinder(ctx context.Context, binder Binder) error {
+func (g *graph) runWithBinder(
+	ctx context.Context,
+	binder Binder,
+	skipped map[*graphNode]struct{},
+) error {
 	rs := &runState{
 		Binder:  binder,
 		signals: map[string]chan struct{}{},
+		skipped: map[string]struct{}{},
 	}
 	for _, gn := range g.nodes {
 		rs.signals[gn.id] = make(chan struct{})
 	}
+	for gn := range skipped {
+		rs.skipped[gn.id] = struct{}{}
+	}
 
 	// errgroup always cancels the derived context before returning from Wait(), so the select below
 	// must listen to the parent context's Done() channel.
 	eg, egCtx := errgroup.WithContext(ctx)
 
 	for _, gn := range g.nodes {
+		if _, isSkipped := skipped[gn]; isSkipped {
+			continue
+		}
 		eg.Go(gn.runFunc(egCtx, rs))
 	}
 
@@ -371,7 +394,8 @@ func (g *graph) AsTask(exposeKeys ...ID) (Task, error) {
 			exposeKeys: exposeSet,
 		}
 
-		if err := g.runWithBinder(ctx, gtb); err != nil {
+		// Overrides apply only to the graph they are given to, not to any graphs it runs as tasks.
+		if err := g.runWithBinder(ctx, gtb, nil); err != nil {
 			return nil, err
 		}
 

internal/overridehook/hook.go

@@ -0,0 +1,26 @@
+// Package overridehook lets taskgraphtest reach the part of the graph engine which replaces tasks.
+//
+// Replacing tasks is meant for tests, so it is not part of the taskgraph package's API. Go only lets
+// packages within this module import an internal package, so this is how taskgraphtest, and nothing
+// outside the module, can get at it. The taskgraph package sets the variables when it is
+// initialised.
+//
+// The functions take the arguments typed as any, as this package can't import taskgraph, which
+// imports it; taskgraphtest knows the types and converts.
+package overridehook
+
+import "context"
+
+var (
+	// RunWithOverrides runs a graph made by taskgraph.New, as Graph.Run does, with some of its tasks
+	// replaced by the override bindings, and returns a taskgraph.Binder. If ignoreUnprovided is set,
+	// overrides for keys which the graph doesn't provide are left out, rather than being an error.
+	RunWithOverrides func(g any, ctx context.Context, inputs, overrides []any, ignoreUnprovided bool) (any, error)
+
+	// CheckWithOverrides is Graph.Check, with overrides.
+	CheckWithOverrides func(g any, inputs, overrides []any) error
+
+	// Provides reports whether some task in a graph made by taskgraph.New provides the key (a
+	// taskgraph.ID).
+	Provides func(g any, id any) bool
+)