Add Override, to replace the task which provides a key

README.md

@@ -27,6 +27,42 @@ 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. `Override` marks an input as
+replacing the task which provides its key:
+
+```go
+graph.Run(ctx, inputs..., tg.Override(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.
+
+* An ordinary 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 ordinary 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. `Graph.Check` finds these
+  without running the graph.
+* The override may be absent: `tg.Override(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 from `Run`, 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 passed to. 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.
+
+`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

@@ -54,11 +54,14 @@
 	//
 	// This is intended to be run by a genrule at build time to assert that all keys required by tasks
 	// in the graph are provided either as an input or by some other task in the graph. It also checks
-	// that there are no duplicate inputs.
+	// that there are no duplicate inputs, and that any Overrides are valid (see Override).
 	Check(inputs ...Binding) error
 
 	// Run executes the task graph with the given inputs, returning a Binder containing the bound
-	// values from all tasks (but not any of the input bindings).
+	// values from all tasks (but not any of the input bindings, other than those given as Overrides,
+	// which stand in for the results of the tasks they replace).
+	//
+	// Inputs given using Override replace the tasks which provide their keys, which are not run.
 	//
 	// It is advisable to set a timeout on the passed context, although it is up to the individual
 	// tasks to listen for context cancellation.
@@ -90,9 +93,16 @@
 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 +137,8 @@
 
 const (
 	traceTaskgraphAbsentKeysPrefix = "taskgraph.absent_keys."
+	traceTaskgraphOverriddenKeys   = "taskgraph.overridden_keys"
+	traceTaskgraphSkippedTasks     = "taskgraph.skipped_tasks"
 )
 
 // Execute the task against the binder provided in the runState.
@@ -262,29 +274,9 @@
 	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...)
 	return err
 }
 
@@ -299,20 +291,33 @@
 		executionLatency.WithLabelValues(g.name, result).
 			Observe(float64(time.Since(startTime) / time.Millisecond))
 	}()
-	base, err := g.buildInputBinder(inputs...)
+	in, err := g.buildInputs(inputs...)
 	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 +327,27 @@
 
 // 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 +383,8 @@
 			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
 		}
 

override.go

@@ -0,0 +1,185 @@
+package taskgraph
+
+import (
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+
+	set "github.com/deckarep/golang-set/v2"
+)
+
+var (
+	// ErrUnknownOverride is returned from Graph.Run() and Graph.Check() if an Override is given for a
+	// key which no task in the graph provides. This is almost always a stale override, left over from
+	// a task which has since been renamed or removed.
+	ErrUnknownOverride = errors.New("override of a key which no task provides")
+
+	// ErrPartialOverride is returned from Graph.Run() and Graph.Check() if Overrides cover some, but
+	// not all, of the keys which a task provides. The task would still have to run, with all of its
+	// effects, only to have some of its results discarded, which is very unlikely to be what was
+	// intended. Override all of a task's keys, or none of them.
+	ErrPartialOverride = errors.New("override of only some of the keys which a task provides")
+)
+
+// An overrideBinding is a Binding which has been marked, using Override, as replacing the task which
+// provides its key.
+type overrideBinding struct {
+	Binding
+}
+
+// Override marks an input Binding as replacing whichever task provides its key, rather than being
+// an ordinary input. It is intended for tests, to substitute the dependencies of a graph which its
+// tasks would otherwise create for themselves:
+//
+//	graph.Run(ctx, inputs..., tg.Override(keyClient.Bind(fakeClient)))
+//
+// The task providing the key is not run. If the task provides other keys as well, they must all be
+// overridden too (see ErrPartialOverride), so that the real task's effects can never happen
+// unnoticed. Tasks which have no other purpose than to produce the inputs of a task which is not run
+// are still run, as are tasks which depend on the overridden key; they see the override as if the
+// task which it replaces had produced it. The override is included in the Binder returned from
+// Run, as the replaced task's results would have been.
+//
+// The Binding may be absent, which replaces the task with a failure:
+//
+//	tg.Override(keyClient.BindError(errUnreachable))
+//
+// Overriding a key which no task provides is an error (ErrUnknownOverride), as is overriding the
+// same key twice, or overriding a key which is also given as an ordinary input. An ordinary input
+// for a key which a task provides remains an error (ErrDuplicateBinding), so a key can only ever be
+// replaced on purpose.
+func Override(b Binding) Binding {
+	if _, ok := b.(overrideBinding); ok {
+		return b
+	}
+	return overrideBinding{Binding: b}
+}
+
+// runInputs is everything a graph is given to run with, once checked against the graph.
+type runInputs struct {
+	// inputs holds the ordinary inputs.
+	inputs Binder
+	// overrides holds the overriding bindings, unwrapped.
+	overrides []Binding
+	// skipped holds the nodes of the tasks which the overrides replace.
+	skipped map[*graphNode]struct{}
+}
+
+// overriddenKeys returns the IDs of the overridden keys, sorted, for reporting.
+func (ri *runInputs) overriddenKeys() []string {
+	var ids []string
+	for _, o := range ri.overrides {
+		ids = append(ids, o.ID().String())
+	}
+	sort.Strings(ids)
+	return ids
+}
+
+// skippedTaskNames returns the names of the tasks which are not run, sorted, for reporting.
+func (ri *runInputs) skippedTaskNames() []string {
+	var names []string
+	for node := range ri.skipped {
+		names = append(names, node.task.Name())
+	}
+	sort.Strings(names)
+	return names
+}
+
+// buildInputs checks the given inputs against the graph, working out which tasks, if any, they
+// replace.
+func (g *graph) buildInputs(inputs ...Binding) (*runInputs, error) {
+	var ordinary, overrides []Binding
+	for _, input := range inputs {
+		if o, ok := input.(overrideBinding); ok {
+			overrides = append(overrides, o.Binding)
+		} else {
+			ordinary = append(ordinary, input)
+		}
+	}
+
+	b := NewBinder()
+	if err := b.Store(ordinary...); err != nil {
+		return nil, wrapStackErrorf("duplicate input: %w", err)
+	}
+
+	result := &runInputs{inputs: b, overrides: overrides, skipped: map[*graphNode]struct{}{}}
+	if len(overrides) > 0 {
+		if err := g.checkOverrides(b, result); err != nil {
+			return nil, err
+		}
+	}
+
+	// The dependencies which have to be given as inputs are those of the tasks which will run, which
+	// are fewer if some are replaced.
+	required := g.allDependencies.Difference(g.allProvided)
+	if len(result.skipped) > 0 {
+		required = set.NewSet[ID]()
+		for _, node := range g.nodes {
+			if _, skipped := result.skipped[node]; !skipped {
+				required.Append(node.task.Depends()...)
+			}
+		}
+		required = required.Difference(g.allProvided)
+	}
+	var missingInputs []string
+	for requiredInput := range required.Iter() {
+		if !b.Has(requiredInput) {
+			missingInputs = append(missingInputs, requiredInput.String())
+		}
+	}
+	if len(missingInputs) > 0 {
+		sort.Strings(missingInputs)
+		return nil, wrapStackErrorf("%w: %s", ErrMissingInputs, strings.Join(missingInputs, ", "))
+	}
+
+	return result, nil
+}
+
+// checkOverrides checks the overrides, filling in which nodes they replace.
+func (g *graph) checkOverrides(ordinary Binder, result *runInputs) error {
+	overridden := set.NewSet[ID]()
+	for _, o := range result.overrides {
+		if ordinary.Has(o.ID()) || overridden.Contains(o.ID()) {
+			return wrapStackErrorf("duplicate input: %w: %q", ErrDuplicateBinding, o.ID())
+		}
+		overridden.Add(o.ID())
+	}
+
+	var unknown []string
+	for id := range overridden.Difference(g.allProvided).Iter() {
+		unknown = append(unknown, id.String())
+	}
+	if len(unknown) > 0 {
+		sort.Strings(unknown)
+		return wrapStackErrorf("%w: %s", ErrUnknownOverride, strings.Join(unknown, ", "))
+	}
+
+	var partial []string
+	for _, node := range g.nodes {
+		provides := node.task.Provides()
+		var covered, uncovered []string
+		for _, id := range provides {
+			if overridden.Contains(id) {
+				covered = append(covered, id.String())
+			} else {
+				uncovered = append(uncovered, id.String())
+			}
+		}
+		switch {
+		case len(covered) == 0:
+		case len(uncovered) == 0:
+			result.skipped[node] = struct{}{}
+		default:
+			partial = append(partial, fmt.Sprintf(
+				"task %s provides [%s, %s], of which only [%s] are overridden",
+				node.task.Name(), strings.Join(covered, ", "), strings.Join(uncovered, ", "), strings.Join(covered, ", "),
+			))
+		}
+	}
+	if len(partial) > 0 {
+		sort.Strings(partial)
+		return wrapStackErrorf("%w: %s", ErrPartialOverride, strings.Join(partial, "; "))
+	}
+	return nil
+}