AroundNode supports simpler signatures too

Author: onsi

core_dsl.go

@@ -330,10 +330,10 @@ func RunSpecs(t GinkgoTestingT, description string, args ...any) bool {
 	return passed
 }
 
-func extractSuiteConfiguration(args []any) (Labels, SemVerConstraints, internal.AroundNodes) {
+func extractSuiteConfiguration(args []any) (Labels, SemVerConstraints, types.AroundNodes) {
 	suiteLabels := Labels{}
 	suiteSemVerConstraints := SemVerConstraints{}
-	aroundNodes := internal.AroundNodes{}
+	aroundNodes := types.AroundNodes{}
 	configErrors := []error{}
 	for _, arg := range args {
 		switch arg := arg.(type) {
@@ -345,7 +345,7 @@ func extractSuiteConfiguration(args []any) (Labels, SemVerConstraints, internal.
 			suiteLabels = append(suiteLabels, arg...)
 		case SemVerConstraints:
 			suiteSemVerConstraints = append(suiteSemVerConstraints, arg...)
-		case internal.AroundNode:
+		case types.AroundNodeDecorator:
 			aroundNodes = append(aroundNodes, arg)
 		default:
 			configErrors = append(configErrors, types.GinkgoErrors.UnknownTypePassedToRunSpecs(arg))

decorator_dsl.go

@@ -2,6 +2,7 @@ package ginkgo
 
 import (
 	"github.com/onsi/ginkgo/v2/internal"
+	"github.com/onsi/ginkgo/v2/types"
 )
 
 /*
@@ -169,4 +170,6 @@ For example, each Ginkgo node runs in its own goroutine.  Some linux namespace c
 
 WIP
 */
-var AroundNode = internal.BuildAroundNode
+func AroundNode[F types.AroundNodeAllowedFuncs](f F) types.AroundNodeDecorator {
+	return types.AroundNode(f, types.NewCodeLocation(1))
+}

dsl/decorators/decorators_dsl.go

@@ -14,6 +14,7 @@ package decorators
 
 import (
 	"github.com/onsi/ginkgo/v2"
+	"github.com/onsi/ginkgo/v2/types"
 )
 
 type Offset = ginkgo.Offset
@@ -37,4 +38,7 @@ const SuppressProgressReporting = ginkgo.SuppressProgressReporting
 
 var Label = ginkgo.Label
 var SemVerConstraint = ginkgo.SemVerConstraint
-var AroundNode = ginkgo.AroundNode
+
+func AroundNode[F types.AroundNodeAllowedFuncs](f F) types.AroundNodeDecorator {
+	return types.AroundNode(f, types.NewCodeLocation(1))
+}

internal/around_node.go

@@ -9,24 +9,21 @@ func ComputeAroundNodes(specs Specs) Specs {
 	for _, spec := range specs {
 		nodes := Nodes{}
 		currentNestingLevel := 0
-		aroundNodes := []AroundNode{}
+		aroundNodes := types.AroundNodes{}
 		nestingLevelIndices := []int{}
 		for _, node := range spec.Nodes {
 			switch node.NodeType {
 			case types.NodeTypeContainer:
 				currentNestingLevel = node.NestingLevel + 1
 				nestingLevelIndices = append(nestingLevelIndices, len(aroundNodes))
-				aroundNodes = append(aroundNodes, node.AroundNodes...)
+				aroundNodes = aroundNodes.Append(node.AroundNodes...)
 				nodes = append(nodes, node)
 			default:
 				if currentNestingLevel > node.NestingLevel {
 					currentNestingLevel = node.NestingLevel
 					aroundNodes = aroundNodes[:nestingLevelIndices[currentNestingLevel]]
 				}
-				nodeAroundNodes := []AroundNode{}
-				nodeAroundNodes = append(nodeAroundNodes, aroundNodes...)
-				nodeAroundNodes = append(nodeAroundNodes, node.AroundNodes...)
-				node.AroundNodes = nodeAroundNodes
+				node.AroundNodes = types.AroundNodes{}.Append(aroundNodes...).Append(node.AroundNodes...)
 				nodes = append(nodes, node)
 			}
 		}

internal/internal_integration/around_node_test.go

@@ -5,15 +5,14 @@ import (
 	"time"
 
 	. "github.com/onsi/ginkgo/v2"
-	"github.com/onsi/ginkgo/v2/internal"
 	. "github.com/onsi/ginkgo/v2/internal/test_helpers"
+	"github.com/onsi/ginkgo/v2/types"
 	. "github.com/onsi/gomega"
 )
 
-func AN(run string) internal.AroundNode {
-	return AroundNode(func(ctx context.Context, body func(ctx context.Context)) {
+func AN(run string) types.AroundNodeDecorator {
+	return AroundNode(func() {
 		rt.Run(run)
-		body(ctx)
 	})
 }
 
@@ -133,10 +132,9 @@ var _ = Describe("The AroundNode decorator", func() {
 					rt.Run("SBS-all")
 					Ω(ctx.Value("wrapped")).Should(Equal("value"))
 					Ω(data).Should(Equal([]byte("data")))
-				}, AroundNode(func(ctx context.Context, body func(ctx context.Context)) {
+				}, AroundNode(func(ctx context.Context) context.Context {
 					rt.Run("SBS-around")
-					newCtx = context.WithValue(ctx, "wrapped", "value")
-					body(newCtx)
+					return context.WithValue(ctx, "wrapped", "value")
 				}))
 				It("runs", rt.T("A"))
 			}, AN("suite-around-1"))

internal/internal_integration/internal_integration_suite_test.go

@@ -88,7 +88,7 @@ func SetUpForParallel(parallelTotal int) {
 	conf.ParallelHost = server.Address()
 }
 
-func RunFixture(description string, callback func(), aroundNodes ...internal.AroundNode) (bool, bool) {
+func RunFixture(description string, callback func(), aroundNodes ...types.AroundNodeDecorator) (bool, bool) {
 	suite := internal.NewSuite()
 	var success, hasProgrammaticFocus bool
 	WithSuite(suite, func() {

internal/node.go

@@ -61,7 +61,7 @@ type Node struct {
 	NodeTimeout             time.Duration
 	SpecTimeout             time.Duration
 	GracePeriod             time.Duration
-	AroundNodes             AroundNodes
+	AroundNodes             types.AroundNodes
 
 	NodeIDWhereCleanupWasGenerated uint
 }
@@ -87,51 +87,24 @@ type FlakeAttempts uint
 type MustPassRepeatedly uint
 type Offset uint
 type Done chan<- any // Deprecated Done Channel for asynchronous testing
-type Labels []string
-type SemVerConstraints []string
 type PollProgressInterval time.Duration
 type PollProgressAfter time.Duration
 type NodeTimeout time.Duration
 type SpecTimeout time.Duration
 type GracePeriod time.Duration
 
-func BuildAroundNode(f AroundNodeFunc) AroundNode {
-	return AroundNode{
-		Body:         f,
-		CodeLocation: types.NewCodeLocation(1),
-	}
-}
-
-type AroundNodeFunc func(ctx context.Context, body func(ctx context.Context))
-
-type AroundNode struct {
-	Body         AroundNodeFunc
-	CodeLocation types.CodeLocation
-}
-
-type AroundNodes []AroundNode
+type Labels []string
 
 func (l Labels) MatchesLabelFilter(query string) bool {
 	return types.MustParseLabelFilter(query)(l)
 }
 
+type SemVerConstraints []string
+
 func (svc SemVerConstraints) MatchesSemVerFilter(version string) bool {
 	return types.MustParseSemVerFilter(version)(svc)
 }
 
-func (an AroundNodes) Clone() AroundNodes {
-	out := make(AroundNodes, len(an))
-	copy(out, an)
-	return out
-}
-
-func (an AroundNodes) Append(other ...AroundNode) AroundNodes {
-	out := make(AroundNodes, len(an)+len(other))
-	copy(out, an)
-	copy(out[len(an):], other)
-	return out
-}
-
 func unionOf[S ~[]E, E comparable](slices ...S) S {
 	out := S{}
 	seen := map[E]bool{}
@@ -207,7 +180,7 @@ func isDecoration(arg any) bool {
 		return true
 	case t == reflect.TypeOf(GracePeriod(0)):
 		return true
-	case t == reflect.TypeOf(AroundNode{}):
+	case t == reflect.TypeOf(types.AroundNodeDecorator{}):
 		return true
 	case t.Kind() == reflect.Slice && isSliceOfDecorations(arg):
 		return true
@@ -349,8 +322,8 @@ func NewNode(deprecationTracker *types.DeprecationTracker, nodeType types.NodeTy
 			if nodeType.Is(types.NodeTypeContainer) {
 				appendError(types.GinkgoErrors.InvalidDecoratorForNodeType(node.CodeLocation, nodeType, "GracePeriod"))
 			}
-		case t == reflect.TypeOf(AroundNode{}):
-			node.AroundNodes = append(node.AroundNodes, arg.(AroundNode))
+		case t == reflect.TypeOf(types.AroundNodeDecorator{}):
+			node.AroundNodes = append(node.AroundNodes, arg.(types.AroundNodeDecorator))
 		case t == reflect.TypeOf(Labels{}):
 			if !nodeType.Is(types.NodeTypesForContainerAndIt) {
 				appendError(types.GinkgoErrors.InvalidDecoratorForNodeType(node.CodeLocation, nodeType, "Label"))

internal/node_test.go

@@ -1228,50 +1228,6 @@ var _ = Describe("Node", func() {
 	})
 })
 
-var _ = Describe("AroundNodes", func() {
-	Describe("building an AroundNode", func() {
-		It("captures the code location", func() {
-			called := false
-			cl := types.NewCodeLocation(0)
-			an := internal.BuildAroundNode(func(ctx context.Context, body func(context.Context)) {
-				called = true
-			})
-			Ω(an).ShouldNot(BeZero())
-			cl.LineNumber += 1
-			Ω(an.CodeLocation).Should(Equal(cl))
-			an.Body(nil, nil)
-			Ω(called).Should(BeTrue())
-		})
-	})
-
-	Describe("Clone", func() {
-		It("clones the slice", func() {
-			an1 := internal.BuildAroundNode(nil)
-			an2 := internal.BuildAroundNode(nil)
-			an3 := internal.BuildAroundNode(nil)
-			original := internal.AroundNodes{an1, an2, an3}
-			clone := original.Clone()
-			Ω(original).Should(Equal(clone))
-			an4 := internal.BuildAroundNode(nil)
-			clone[2] = an4
-			Ω(original).Should(Equal(internal.AroundNodes{an1, an2, an3}))
-			Ω(clone).Should(Equal(internal.AroundNodes{an1, an2, an4}))
-		})
-	})
-
-	Describe("Append", func() {
-		It("appends the node to the slice", func() {
-			an1 := internal.BuildAroundNode(nil)
-			an2 := internal.BuildAroundNode(nil)
-			an3 := internal.BuildAroundNode(nil)
-			nodes := internal.AroundNodes{an1, an2}
-			newNodes := nodes.Append(an3)
-			Ω(nodes).Should(Equal(internal.AroundNodes{an1, an2}))
-			Ω(newNodes).Should(Equal(internal.AroundNodes{an1, an2, an3}))
-		})
-	})
-})
-
 var _ = Describe("Nodes", func() {
 	Describe("Clone", func() {
 		var n1, n2, n3, n4 Node

internal/suite.go

@@ -32,7 +32,7 @@ type Suite struct {
 
 	suiteNodes   Nodes
 	cleanupNodes Nodes
-	aroundNodes  AroundNodes
+	aroundNodes  types.AroundNodes
 
 	failer            *Failer
 	reporter          reporters.Reporter
@@ -106,7 +106,7 @@ func (suite *Suite) BuildTree() error {
 	return nil
 }
 
-func (suite *Suite) Run(description string, suiteLabels Labels, suiteSemVerConstraints SemVerConstraints, suiteAroundNodes AroundNodes, suitePath string, failer *Failer, reporter reporters.Reporter, writer WriterInterface, outputInterceptor OutputInterceptor, interruptHandler interrupt_handler.InterruptHandlerInterface, client parallel_support.Client, progressSignalRegistrar ProgressSignalRegistrar, suiteConfig types.SuiteConfig) (bool, bool) {
+func (suite *Suite) Run(description string, suiteLabels Labels, suiteSemVerConstraints SemVerConstraints, suiteAroundNodes types.AroundNodes, suitePath string, failer *Failer, reporter reporters.Reporter, writer WriterInterface, outputInterceptor OutputInterceptor, interruptHandler interrupt_handler.InterruptHandlerInterface, client parallel_support.Client, progressSignalRegistrar ProgressSignalRegistrar, suiteConfig types.SuiteConfig) (bool, bool) {
 	if suite.phase != PhaseBuildTree {
 		panic("cannot run before building the tree = call suite.BuildTree() first")
 	}
@@ -263,7 +263,7 @@ func (suite *Suite) pushCleanupNode(node Node) error {
 
 	node.NodeIDWhereCleanupWasGenerated = suite.currentNode.ID
 	node.NestingLevel = suite.currentNode.NestingLevel
-	node.AroundNodes = AroundNodes{}.Append(suite.currentNode.AroundNodes...).Append(node.AroundNodes...)
+	node.AroundNodes = types.AroundNodes{}.Append(suite.currentNode.AroundNodes...).Append(node.AroundNodes...)
 	suite.selectiveLock.Lock()
 	suite.cleanupNodes = append(suite.cleanupNodes, node)
 	suite.selectiveLock.Unlock()
@@ -897,7 +897,7 @@ func (suite *Suite) runNode(node Node, specDeadline time.Time, text string) (typ
 			failureC <- failureFromRun
 		}()
 
-		aroundNodes := AroundNodes{}.Append(suite.aroundNodes...).Append(node.AroundNodes...)
+		aroundNodes := types.AroundNodes{}.Append(suite.aroundNodes...).Append(node.AroundNodes...)
 		if len(aroundNodes) > 0 {
 			i := 0
 			var f func(context.Context)

types/around_node.go

@@ -0,0 +1,56 @@
+package types
+
+import (
+	"context"
+)
+
+type AroundNodeAllowedFuncs interface {
+	~func(context.Context, func(context.Context)) | ~func(context.Context) context.Context | ~func()
+}
+type AroundNodeFunc func(ctx context.Context, body func(ctx context.Context))
+
+func AroundNode[F AroundNodeAllowedFuncs](f F, cl CodeLocation) AroundNodeDecorator {
+	if f == nil {
+		panic("BuildAroundNode cannot be called with a nil function.")
+	}
+	var aroundNodeFunc func(context.Context, func(context.Context))
+	switch x := any(f).(type) {
+	case func(context.Context, func(context.Context)):
+		aroundNodeFunc = x
+	case func(context.Context) context.Context:
+		aroundNodeFunc = func(ctx context.Context, body func(context.Context)) {
+			ctx = x(ctx)
+			body(ctx)
+		}
+	case func():
+		aroundNodeFunc = func(ctx context.Context, body func(context.Context)) {
+			x()
+			body(ctx)
+		}
+	}
+
+	return AroundNodeDecorator{
+		Body:         aroundNodeFunc,
+		CodeLocation: cl,
+	}
+}
+
+type AroundNodeDecorator struct {
+	Body         AroundNodeFunc
+	CodeLocation CodeLocation
+}
+
+type AroundNodes []AroundNodeDecorator
+
+func (an AroundNodes) Clone() AroundNodes {
+	out := make(AroundNodes, len(an))
+	copy(out, an)
+	return out
+}
+
+func (an AroundNodes) Append(other ...AroundNodeDecorator) AroundNodes {
+	out := make(AroundNodes, len(an)+len(other))
+	copy(out, an)
+	copy(out[len(an):], other)
+	return out
+}