Repository: kubernetes-sigs/kwok (go)

Issue:
Add a subcommand for debugging the Stage
### What would you like to be added?

```
# Output matched stages with the resource
kwokctl tool stage match <resource file> <stage file...>
```

```
# Output stage's actions with the resource
kwokctl tool stage action <resource file> <stage file> 
```

### Why is this needed?

In the future, we will try to enhance kwok's ability to simulate various behaviors of resources, and we need an easy way to debug it

Pull request description:
Add test for stages
<!--  Thanks for sending a pull request!  Here are some tips for you:

1. If this is your first time, please read our contributor guidelines: https://git.k8s.io/community/contributors/guide/first-contribution.md#your-first-contribution and developer guide https://git.k8s.io/community/contributors/devel/development.md#development-guide
2. Please label this pull request according to what type of issue you are addressing, especially if this is a release targeted pull request. For reference on required PR/issue labels, read here:
https://git.k8s.io/community/contributors/devel/sig-release/release.md#issuepr-kind-label
3. Ensure you have added or ran the appropriate tests for your PR: https://git.k8s.io/community/contributors/devel/sig-testing/testing.md
4. If you want *faster* PR reviews, read how: https://git.k8s.io/community/contributors/guide/pull-requests.md#best-practices-for-faster-reviews
5. If the PR is unfinished, see how to mark it: https://git.k8s.io/community/contributors/guide/pull-requests.md#marking-unfinished-pull-requests
-->

#### What type of PR is this?

<!--
Add one of the following kinds:
/kind bug
/kind cleanup
/kind documentation
/kind feature

Optionally add one or more of the following kinds if applicable:
/kind api-change
/kind deprecation
/kind failing-test
/kind flake
/kind regression
-->

/kind feature

#### What this PR does / why we need it:

#### Which issue(s) this PR fixes:
<!--
*Automatically closes linked issue when PR is merged.
Usage: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
_If PR is about `failing-tests or flakes`, please post the related issues/tests in a comment and do not use `Fixes`_*
-->
Fixes #1027 
Based on #1067

#### Special notes for your reviewer:

#### Does this PR introduce a user-facing change?
<!--
If no, just write "NONE" in the release-note block below.
If yes, a release note is required:
Enter your extended release note in the block below. If the PR requires additional action from users switching to the new release, include the string "action required".

For more information on release notes see: https://git.k8s.io/community/contributors/guide/release-notes.md
-->
```release-note
NONE
```

#### Additional documentation e.g., KEPs (Kubernetes Enhancement Proposals), usage docs, etc.:

<!--
This section can be blank if this pull request does not require a release note.

When adding links which point to resources within git repositories, like
KEPs or supporting documentation, please reference a specific commit and avoid
linking directly to the master branch. This ensures that links reference a
specific point in time, rather than a document that may change over time.

See here for guidance on getting permanent links to files: https://help.github.com/en/articles/getting-permanent-links-to-files

Please use the following format for linking documentation:
- [KEP]: <link>
- [Usage]: <link>
- [Other doc]: <link>
-->
```docs
NONE
```


Patch:
diff --git a/hack/update-all.sh b/hack/update-all.sh
index 03d10316..91dc5d9a 100755
--- a/hack/update-all.sh
+++ b/hack/update-all.sh
@@ -78,6 +78,11 @@ if [[ "${UPDATE_SPELLING:-true}" == "true" ]]; then
   "${ROOT_DIR}"/hack/update-spelling.sh || failed+=(spelling)
 fi
 
+if [[ "${UPDATE_STAGES:-true}" == "true" ]]; then
+  echo "[*] Update stages..."
+  "${ROOT_DIR}"/hack/update-stages.sh || failed+=(stages)
+fi
+
 if [[ "${#failed[@]}" != 0 ]]; then
   echo "Update failed for: ${failed[*]}"
   exit 1
diff --git a/hack/update-stages.sh b/hack/update-stages.sh
new file mode 100755
index 00000000..2c198a29
--- /dev/null
+++ b/hack/update-stages.sh
@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+# Copyright 2024 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+DIR="$(dirname "${BASH_SOURCE[0]}")"
+
+ROOT_DIR="$(realpath "${DIR}/..")"
+
+function update() {
+  mapfile -t files < <(
+    find . \
+      -iname "*.input.yaml"
+  )
+
+  for file in "${files[@]}"; do
+    update_stage "${file}"
+  done
+}
+
+function update_stage() {
+  local file="${1}"
+  local stages=()
+  local rel_path
+  local base_dir
+  local from
+  base_dir="$(dirname "${file}")"
+
+  from="$(grep '# @Stage: ' "${file}" | awk '{print $3}')"
+  for line in ${from}; do
+    rel_path="${line}"
+    stages+=("${base_dir}/${rel_path}")
+  done
+
+  if [[ ${#stages[@]} -eq 0 ]]; then
+    return
+  fi
+
+  go run ./hack/test_stage "${file}" "${stages[@]}" >"${file%.input.yaml}.output.yaml"
+}
+
+cd "${ROOT_DIR}" && update
diff --git a/hack/verify-all.sh b/hack/verify-all.sh
index d68ee57d..bf774b34 100755
--- a/hack/verify-all.sh
+++ b/hack/verify-all.sh
@@ -93,6 +93,11 @@ if [[ "${VERIFY_SPELLING:-true}" == "true" ]]; then
   "${ROOT_DIR}"/hack/verify-spelling.sh || failed+=(spelling)
 fi
 
+if [[ "${VERIFY_STAGES:-true}" == "true" ]]; then
+  echo "[*] Verifying stages..."
+  "${ROOT_DIR}"/hack/verify-stages.sh || failed+=(stages)
+fi
+
 # exit based on verify scripts
 if [[ "${#failed[@]}" != 0 ]]; then
   echo "Verify failed for: ${failed[*]}"
diff --git a/hack/verify-stages.sh b/hack/verify-stages.sh
new file mode 100755
index 00000000..1ba27219
--- /dev/null
+++ b/hack/verify-stages.sh
@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+# Copyright 2024 The Kubernetes Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+set -o errexit
+set -o nounset
+set -o pipefail
+
+DIR="$(dirname "${BASH_SOURCE[0]}")"
+
+ROOT_DIR="$(realpath "${DIR}/..")"
+
+function check() {
+  rm -rf "${ROOT_DIR}"/kustomize/stage/**/testdata/*.output.yaml
+  "${ROOT_DIR}"/hack/update-stages.sh
+  git --no-pager diff --exit-code
+}
+
+cd "${ROOT_DIR}" && check
diff --git a/pkg/config/config.go b/pkg/config/config.go
index 84b8298f..02de27a1 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -27,6 +27,7 @@ import (
 	"sort"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
 	"k8s.io/apimachinery/pkg/runtime/schema"
 
 	configv1alpha1 "sigs.k8s.io/kwok/pkg/apis/config/v1alpha1"
@@ -328,6 +329,27 @@ func Load(ctx context.Context, src ...string) ([]InternalObject, error) {
 	return objs, nil
 }
 
+// LoadUnstructured loads the given path into the context.
+func LoadUnstructured(src ...string) ([]InternalObject, error) {
+	raws, err := loadRawMessages(src)
+	if err != nil {
+		return nil, err
+	}
+
+	objs := []InternalObject{}
+	for _, raw := range raws {
+		obj := unstructured.Unstructured{}
+		err := json.Unmarshal(raw, &obj)
+		if err != nil {
+			return nil, err
+		}
+
+		objs = append(objs, &obj)
+	}
+
+	return objs, nil
+}
+
 // Save saves the given objects to the given path.
 func Save(ctx context.Context, dist string, objs []InternalObject) error {
 	dist = path.Clean(dist)
diff --git a/pkg/kwok/controllers/controller.go b/pkg/kwok/controllers/controller.go
index abb0785c..1b39c798 100644
--- a/pkg/kwok/controllers/controller.go
+++ b/pkg/kwok/controllers/controller.go
@@ -18,11 +18,8 @@ package controllers
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"os"
-	"strconv"
-	"strings"
 	"time"
 
 	coordinationv1 "k8s.io/api/coordination/v1"
@@ -45,7 +42,6 @@ import (
 	"sigs.k8s.io/kwok/pkg/apis/v1alpha1"
 	"sigs.k8s.io/kwok/pkg/client/clientset/versioned"
 	"sigs.k8s.io/kwok/pkg/config/resources"
-	"sigs.k8s.io/kwok/pkg/consts"
 	"sigs.k8s.io/kwok/pkg/log"
 	"sigs.k8s.io/kwok/pkg/utils/client"
 	"sigs.k8s.io/kwok/pkg/utils/gotpl"
@@ -54,50 +50,9 @@ import (
 	"sigs.k8s.io/kwok/pkg/utils/patch"
 	"sigs.k8s.io/kwok/pkg/utils/queue"
 	"sigs.k8s.io/kwok/pkg/utils/slices"
-	"sigs.k8s.io/kwok/pkg/utils/yaml"
 )
 
 var (
-	startTime = time.Now().Format(time.RFC3339Nano)
-
-	defaultFuncMap = gotpl.FuncMap{
-		"Quote": func(s any) string {
-			data, err := json.Marshal(s)
-			if err != nil {
-				return strconv.Quote(fmt.Sprint(s))
-			}
-			if len(data) == 0 {
-				return `""`
-			}
-			if data[0] == '"' {
-				return string(data)
-			}
-			return strconv.Quote(string(data))
-		},
-		"Now": func() string {
-			return time.Now().Format(time.RFC3339Nano)
-		},
-		"StartTime": func() string {
-			return startTime
-		},
-		"YAML": func(s interface{}, indent ...int) (string, error) {
-			d, err := yaml.Marshal(s)
-			if err != nil {
-				return "", err
-			}
-
-			data := string(d)
-			if len(indent) == 1 && indent[0] > 0 {
-				pad := strings.Repeat(" ", indent[0]*2)
-				data = strings.ReplaceAll("\n"+data, "\n", "\n"+pad)
-			}
-			return data, nil
-		},
-		"Version": func() string {
-			return consts.Version
-		},
-	}
-
 	nodeKind = corev1.SchemeGroupVersion.WithKind("Node")
 )
 
@@ -170,6 +125,7 @@ type Config struct {
 	ID                                    string
 	EnableMetrics                         bool
 	EnablePodCache                        bool
+	FuncMap                               gotpl.FuncMap
 }
 
 func (c Config) validate() error {
@@ -471,7 +427,7 @@ func (c *Controller) initNodeController(ctx context.Context, lifecycle resources
 		OnNodeUnmanagedFunc:                   c.onNodeUnmanaged,
 		Lifecycle:                             lifecycle,
 		PlayStageParallelism:                  c.conf.NodePlayStageParallelism,
-		FuncMap:                               defaultFuncMap,
+		FuncMap:                               c.conf.FuncMap,
 		Recorder:                              c.recorder,
 		ReadOnlyFunc:                          c.readOnlyFunc,
 		EnableMetrics:                         c.conf.EnableMetrics,
@@ -505,7 +461,7 @@ func (c *Controller) initPodController(ctx context.Context, lifecycle resources.
 
 			return c.nodes.Get(nodeName)
 		},
-		FuncMap:       defaultFuncMap,
+		FuncMap:       c.conf.FuncMap,
 		Recorder:      c.recorder,
 		ReadOnlyFunc:  c.readOnlyFunc,
 		EnableMetrics: c.conf.EnableMetrics,
@@ -558,7 +514,7 @@ func (c *Controller) initStageController(ctx context.Context, ref internalversio
 		DisregardStatusWithLabelSelector:      c.conf.DisregardStatusWithLabelSelector,
 		Lifecycle:                             lifecycle,
 		PlayStageParallelism:                  1,
-		FuncMap:                               defaultFuncMap,
+		FuncMap:                               c.conf.FuncMap,
 		Recorder:                              c.recorder,
 	})
 	if err != nil {
diff --git a/pkg/kwok/controllers/node_controller.go b/pkg/kwok/controllers/node_controller.go
index 370700c8..fb63e974 100644
--- a/pkg/kwok/controllers/node_controller.go
+++ b/pkg/kwok/controllers/node_controller.go
@@ -42,43 +42,6 @@ import (
 	"sigs.k8s.io/kwok/pkg/utils/wait"
 )
 
-var (
-	// https://kubernetes.io/docs/concepts/architecture/nodes/#condition
-	nodeConditions = []corev1.NodeCondition{
-		{
-			Type:    corev1.NodeReady,
-			Status:  corev1.ConditionTrue,
-			Reason:  "KubeletReady",
-			Message: "kubelet is posting ready status",
-		},
-		{
-			Type:    corev1.NodeMemoryPressure,
-			Status:  corev1.ConditionFalse,
-			Reason:  "KubeletHasSufficientMemory",
-			Message: "kubelet has sufficient memory available",
-		},
-		{
-			Type:    corev1.NodeDiskPressure,
-			Status:  corev1.ConditionFalse,
-			Reason:  "KubeletHasNoDiskPressure",
-			Message: "kubelet has no disk pressure",
-		},
-		{
-			Type:    corev1.NodePIDPressure,
-			Status:  corev1.ConditionFalse,
-			Reason:  "KubeletHasSufficientPID",
-			Message: "kubelet has sufficient PID available",
-		},
-		{
-			Type:    corev1.NodeNetworkUnavailable,
-			Status:  corev1.ConditionFalse,
-			Reason:  "RouteCreated",
-			Message: "RouteController created a route",
-		},
-	}
-	nodeConditionsData, _ = expression.ToJSONStandard(nodeConditions)
-)
-
 // NodeController is a fake nodes implementation that can be used to test
 type NodeController struct {
 	clock                                 clock.Clock
@@ -171,9 +134,6 @@ func NewNodeController(conf NodeControllerConfig) (*NodeController, error) {
 		"NodeIP":   c.funcNodeIP,
 		"NodeName": c.funcNodeName,
 		"NodePort": c.funcNodePort,
-		"NodeConditions": func() interface{} {
-			return nodeConditionsData
-		},
 	}, conf.FuncMap)
 	c.renderer = gotpl.NewRenderer(funcMap)
 	return c, nil
diff --git a/pkg/tools/stage/stage.go b/pkg/tools/stage/stage.go
new file mode 100644
index 00000000..c92f5e59
--- /dev/null
+++ b/pkg/tools/stage/stage.go
@@ -0,0 +1,215 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package stage provides a stage tester.
+package stage
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"time"
+
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+
+	"sigs.k8s.io/kwok/pkg/apis/internalversion"
+	"sigs.k8s.io/kwok/pkg/utils/gotpl"
+	"sigs.k8s.io/kwok/pkg/utils/lifecycle"
+	"sigs.k8s.io/kwok/pkg/utils/slices"
+)
+
+// TestingStages tests the stages against the target object
+func TestingStages(ctx context.Context, target any, stages []*internalversion.Stage) (any, error) {
+	testTarget, ok := target.(Obj)
+	if !ok {
+		return nil, fmt.Errorf("expected target to be an object, got %T", target)
+	}
+
+	gvk := testTarget.GetObjectKind().GroupVersionKind()
+
+	want := internalversion.StageResourceRef{
+		APIGroup: gvk.GroupVersion().String(),
+		Kind:     gvk.Kind,
+	}
+
+	meta := map[string]any{
+		"apiGroup": want.APIGroup,
+		"kind":     want.Kind,
+		"name":     testTarget.GetName(),
+	}
+	if ns := testTarget.GetNamespace(); ns != "" {
+		meta["namespace"] = ns
+	}
+
+	stages = slices.Filter(stages, func(stage *internalversion.Stage) bool {
+		return stage.Spec.ResourceRef == want
+	})
+
+	lc, err := lifecycle.NewLifecycle(stages)
+	if err != nil {
+		return nil, err
+	}
+
+	lcstages, err := lc.ListAllPossible(testTarget.GetLabels(), testTarget.GetAnnotations(), testTarget)
+	if err != nil {
+		return nil, err
+	}
+
+	out := []any{}
+	for _, stage := range lcstages {
+		o, err := testingStage(ctx, testTarget, stage)
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, o)
+	}
+
+	meta["stages"] = out
+
+	return meta, nil
+}
+
+func testingStage(ctx context.Context, testTarget Obj, stage *lifecycle.Stage) (any, error) {
+	meta := map[string]any{
+		"stage": stage.Name(),
+	}
+
+	delay, ok := stage.Delay(ctx, stage, time.Now())
+	if ok {
+		meta["delay"] = delay
+	}
+
+	weight := stage.Weight()
+	if weight != 0 {
+		meta["weight"] = weight
+	}
+
+	next := stage.Next()
+
+	if next == nil {
+		meta["next"] = "nil"
+		return meta, nil
+	}
+
+	out := make([]any, 0)
+
+	patch, err := next.Finalizers(testTarget.GetFinalizers())
+	if err != nil {
+		return nil, err
+	}
+
+	if patch != nil {
+		out = append(out, formatPatch(patch))
+	}
+
+	if next.Delete() {
+		out = append(out, map[string]string{
+			"kind": "delete",
+		})
+		meta["next"] = out
+		return meta, nil
+	}
+
+	fm := gotpl.FuncMap{}
+	funcNames := []string{
+		// For node and pod
+		"NodeIP",
+
+		// For node
+		"NodeName",
+		"NodePort",
+
+		// For pod
+		"PodIP",
+		"NodeIPWith",
+		"PodIPWith",
+
+		// Override built-in
+		"Now",
+		"now",
+		"Version",
+	}
+	for _, name := range funcNames {
+		fm[name] = wrapFunction(name)
+	}
+
+	renderer := gotpl.NewRenderer(fm)
+
+	patches, err := next.Patches(testTarget, renderer)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, patch := range patches {
+		out = append(out, formatPatch(patch))
+	}
+
+	if stage.ImmediateNextStage() {
+		out = append(out, map[string]string{
+			"kind": "immediate",
+		})
+	}
+
+	meta["next"] = out
+	return meta, nil
+}
+
+func wrapFunction(name string) func(args ...any) any {
+	return func(args ...any) any {
+		if len(args) == 0 {
+			return fmt.Sprintf("<%s>", name)
+		}
+
+		return fmt.Sprintf("<%s(%s)>", name,
+			strings.Join(
+				slices.Map(args,
+					func(arg any) string {
+						a := fmt.Sprintf("%#v", arg)
+						if a == "" {
+							return `""`
+						}
+						return a
+					},
+				),
+				", ",
+			),
+		)
+	}
+}
+
+func formatPatch(patch *lifecycle.Patch) any {
+	out := map[string]any{
+		"kind": "patch",
+	}
+	out["type"] = patch.Type
+	if patch.Subresource != "" {
+		out["subresource"] = patch.Subresource
+	}
+	out["data"] = json.RawMessage(patch.Data)
+
+	if patch.Impersonation != nil {
+		out["impersonation"] = patch.Impersonation.Username
+	}
+
+	return out
+}
+
+type Obj interface {
+	metav1.Object
+	runtime.Object
+}
diff --git a/pkg/utils/gotpl/funcs.go b/pkg/utils/gotpl/funcs.go
index 73303be4..40c2c081 100644
--- a/pkg/utils/gotpl/funcs.go
+++ b/pkg/utils/gotpl/funcs.go
@@ -17,10 +17,102 @@ limitations under the License.
 package gotpl
 
 import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
 	"github.com/Masterminds/sprig/v3"
+	corev1 "k8s.io/api/core/v1"
+
+	"sigs.k8s.io/kwok/pkg/consts"
+	"sigs.k8s.io/kwok/pkg/utils/expression"
+	"sigs.k8s.io/kwok/pkg/utils/yaml"
 )
 
 var (
 	// genericFuncs is generic template functions.
 	genericFuncs = sprig.TxtFuncMap()
 )
+
+var (
+	startTime = time.Now().Format(time.RFC3339Nano)
+
+	defaultFuncs = FuncMap{
+		"Quote": func(s any) string {
+			data, err := json.Marshal(s)
+			if err != nil {
+				return strconv.Quote(fmt.Sprint(s))
+			}
+			if len(data) == 0 {
+				return `""`
+			}
+			if data[0] == '"' {
+				return string(data)
+			}
+			return strconv.Quote(string(data))
+		},
+		"Now": func() string {
+			return time.Now().Format(time.RFC3339Nano)
+		},
+		"StartTime": func() string {
+			return startTime
+		},
+		"YAML": func(s interface{}, indent ...int) (string, error) {
+			d, err := yaml.Marshal(s)
+			if err != nil {
+				return "", err
+			}
+
+			data := string(d)
+			if len(indent) == 1 && indent[0] > 0 {
+				pad := strings.Repeat(" ", indent[0]*2)
+				data = strings.ReplaceAll("\n"+data, "\n", "\n"+pad)
+			}
+			return data, nil
+		},
+		"Version": func() string {
+			return consts.Version
+		},
+
+		"NodeConditions": func() interface{} {
+			return nodeConditionsData
+		},
+	}
+
+	// https://kubernetes.io/docs/concepts/architecture/nodes/#condition
+	nodeConditions = []corev1.NodeCondition{
+		{
+			Type:    corev1.NodeReady,
+			Status:  corev1.ConditionTrue,
+			Reason:  "KubeletReady",
+			Message: "kubelet is posting ready status",
+		},
+		{
+			Type:    corev1.NodeMemoryPressure,
+			Status:  corev1.ConditionFalse,
+			Reason:  "KubeletHasSufficientMemory",
+			Message: "kubelet has sufficient memory available",
+		},
+		{
+			Type:    corev1.NodeDiskPressure,
+			Status:  corev1.ConditionFalse,
+			Reason:  "KubeletHasNoDiskPressure",
+			Message: "kubelet has no disk pressure",
+		},
+		{
+			Type:    corev1.NodePIDPressure,
+			Status:  corev1.ConditionFalse,
+			Reason:  "KubeletHasSufficientPID",
+			Message: "kubelet has sufficient PID available",
+		},
+		{
+			Type:    corev1.NodeNetworkUnavailable,
+			Status:  corev1.ConditionFalse,
+			Reason:  "RouteCreated",
+			Message: "RouteController created a route",
+		},
+	}
+	nodeConditionsData, _ = expression.ToJSONStandard(nodeConditions)
+)
diff --git a/pkg/utils/gotpl/renderer.go b/pkg/utils/gotpl/renderer.go
index 88ef1f9f..26996a95 100644
--- a/pkg/utils/gotpl/renderer.go
+++ b/pkg/utils/gotpl/renderer.go
@@ -63,6 +63,7 @@ func (r *renderer) render(buf *bytes.Buffer, text string, original interface{})
 		var err error
 		temp, err = template.New("_").
 			Funcs(genericFuncs).
+			Funcs(defaultFuncs).
 			Funcs(r.funcMap).
 			Parse(text)
 		if err != nil {
diff --git a/pkg/utils/lifecycle/lifecycle.go b/pkg/utils/lifecycle/lifecycle.go
index 5eacc8be..16b60e16 100644
--- a/pkg/utils/lifecycle/lifecycle.go
+++ b/pkg/utils/lifecycle/lifecycle.go
@@ -62,6 +62,41 @@ func (s Lifecycle) match(label, annotation labels.Set, data interface{}) ([]*Sta
 	return out, nil
 }
 
+// ListAllPossible returns all possible stages.
+func (s Lifecycle) ListAllPossible(label, annotation labels.Set, data interface{}) ([]*Stage, error) {
+	data, err := expression.ToJSONStandard(data)
+	if err != nil {
+		return nil, err
+	}
+	stages, err := s.match(label, annotation, data)
+	if err != nil {
+		return nil, err
+	}
+	if len(stages) == 0 {
+		return nil, nil
+	}
+	if len(stages) == 1 {
+		return stages, nil
+	}
+
+	totalWeights := 0
+	for _, stage := range stages {
+		totalWeights += stage.weight
+	}
+	if totalWeights == 0 {
+		return stages, nil
+	}
+
+	stagesWithWeights := make([]*Stage, 0, len(stages))
+	for _, stage := range stages {
+		if stage.weight == 0 {
+			continue
+		}
+		stagesWithWeights = append(stagesWithWeights, stage)
+	}
+	return stagesWithWeights, nil
+}
+
 // Match returns matched stage.
 func (s Lifecycle) Match(label, annotation labels.Set, data interface{}) (*Stage, error) {
 	data, err := expression.ToJSONStandard(data)
@@ -256,3 +291,8 @@ func (s *Stage) Name() string {
 func (s *Stage) ImmediateNextStage() bool {
 	return s.immediateNextStage
 }
+
+// Weight returns the weight of the stage.
+func (s *Stage) Weight() int {
+	return s.weight
+}
diff --git a/pkg/utils/yaml/encoder.go b/pkg/utils/yaml/encoder.go
index 1879372e..89ad804e 100644
--- a/pkg/utils/yaml/encoder.go
+++ b/pkg/utils/yaml/encoder.go
@@ -19,8 +19,6 @@ package yaml
 import (
 	"io"
 	"sync/atomic"
-
-	"k8s.io/apimachinery/pkg/runtime"
 )
 
 // Encoder is a YAML encoder.
@@ -39,7 +37,7 @@ func NewEncoder(w io.Writer) *Encoder {
 var separator = []byte("---\n")
 
 // Encode prints the object as YAML.
-func (p *Encoder) Encode(obj runtime.Object) error {
+func (p *Encoder) Encode(obj any) error {
 	count := atomic.AddInt64(&p.printCount, 1)
 	if count > 1 {
 		if _, err := p.w.Write(separator); err != nil {
diff --git a/site/content/en/docs/user/kwok-in-cluster.md b/site/content/en/docs/user/kwok-in-cluster.md
index 7daf0ad8..7a2e38bf 100644
--- a/site/content/en/docs/user/kwok-in-cluster.md
+++ b/site/content/en/docs/user/kwok-in-cluster.md
@@ -19,13 +19,13 @@ KWOK_REPO=kubernetes-sigs/kwok
 KWOK_LATEST_RELEASE=$(curl "https://api.github.com/repos/${KWOK_REPO}/releases/latest" | jq -r '.tag_name')
 ```
 
-## Deploy kwok and set up custom resource definitions (CRDs)
+## Deployment kwok and set up CRDs
 
 ``` bash
 kubectl apply -f "https://github.com/${KWOK_REPO}/releases/download/${KWOK_LATEST_RELEASE}/kwok.yaml"
 ```
 
-## Set up default custom resources (CRs) of stages (required)
+## Set up default CRs of Stages (required)
 
 {{< hint "warning" >}}
 NOTE: This configures the pod/node emulation behavior, if not it will do nothing.
@@ -35,7 +35,7 @@ NOTE: This configures the pod/node emulation behavior, if not it will do nothing
 kubectl apply -f "https://github.com/${KWOK_REPO}/releases/download/${KWOK_LATEST_RELEASE}/stage-fast.yaml"
 ```
 
-## Set up default custom resources (CRs) of resource usage (optional)
+## Set up default CRs of resource usage (optional)
 
 This allows to simulate the resource usage of nodes, pods and containers.
 


Test patch:
diff --git a/hack/test_stage/main.go b/hack/test_stage/main.go
new file mode 100644
index 00000000..caf191ba
--- /dev/null
+++ b/hack/test_stage/main.go
@@ -0,0 +1,81 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+// Package main implements a simple stage tester.
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+
+	"sigs.k8s.io/kwok/pkg/apis/internalversion"
+	"sigs.k8s.io/kwok/pkg/config"
+	"sigs.k8s.io/kwok/pkg/tools/stage"
+	"sigs.k8s.io/kwok/pkg/utils/yaml"
+)
+
+func main() {
+	args := os.Args[1:]
+	if len(args) < 2 {
+		fmt.Fprintf(os.Stderr, "usage: %s <resource file> <stage file...>\n", os.Args[0])
+		os.Exit(1)
+	}
+
+	ctx := context.Background()
+	err := run(ctx, args[0], args[1:])
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "error: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+func run(ctx context.Context, resourceFile string, stageFile []string) error {
+	rio, err := config.LoadUnstructured(resourceFile)
+	if err != nil {
+		return err
+	}
+	if len(rio) != 1 {
+		return fmt.Errorf("expected exactly one resource in %s, got %d", resourceFile, len(rio))
+	}
+
+	sio, err := config.Load(ctx, stageFile...)
+	if err != nil {
+		return err
+	}
+
+	ue := config.FilterWithoutType[*internalversion.Stage](sio)
+	if len(ue) != 0 {
+		return fmt.Errorf("expected only stage, got %d non-stage", len(ue))
+	}
+
+	stages := config.FilterWithType[*internalversion.Stage](sio)
+	if len(stages) == 0 {
+		return fmt.Errorf("expected at least one stage, got 0")
+	}
+
+	out, err := stage.TestingStages(ctx, rio[0], stages)
+	if err != nil {
+		return err
+	}
+
+	err = yaml.NewEncoder(os.Stdout).Encode(out)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
diff --git a/kustomize/stage/node/fast/testdata/node-pending.input.yaml b/kustomize/stage/node/fast/testdata/node-pending.input.yaml
new file mode 100644
index 00000000..2655c60d
--- /dev/null
+++ b/kustomize/stage/node/fast/testdata/node-pending.input.yaml
@@ -0,0 +1,5 @@
+# @Stage: ../node-initialize.yaml
+apiVersion: v1
+kind: Node
+metadata:
+  name: node-pending
diff --git a/kustomize/stage/node/fast/testdata/node-pending.output.yaml b/kustomize/stage/node/fast/testdata/node-pending.output.yaml
new file mode 100644
index 00000000..c822e729
--- /dev/null
+++ b/kustomize/stage/node/fast/testdata/node-pending.output.yaml
@@ -0,0 +1,59 @@
+apiGroup: v1
+kind: Node
+name: node-pending
+stages:
+- next:
+  - data:
+      status:
+        addresses:
+        - address: <NodeIP>
+          type: InternalIP
+        - address: <NodeName>
+          type: Hostname
+        allocatable:
+          cpu: 1k
+          memory: 1Ti
+          pods: 1M
+        capacity:
+          cpu: 1k
+          memory: 1Ti
+          pods: 1M
+        conditions:
+        - lastHeartbeatTime: <Now>
+          lastTransitionTime: <Now>
+          message: kubelet is posting ready status
+          reason: KubeletReady
+          status: "True"
+          type: Ready
+        - lastHeartbeatTime: <Now>
+          lastTransitionTime: <Now>
+          message: kubelet has sufficient memory available
+          reason: KubeletHasSufficientMemory
+          status: "False"
+          type: MemoryPressure
+        - lastHeartbeatTime: <Now>
+          lastTransitionTime: <Now>
+          message: kubelet has no disk pressure
+          reason: KubeletHasNoDiskPressure
+          status: "False"
+          type: DiskPressure
+        - lastHeartbeatTime: <Now>
+          lastTransitionTime: <Now>
+          message: kubelet has sufficient PID available
+          reason: KubeletHasSufficientPID
+          status: "False"
+          type: PIDPressure
+        - lastHeartbeatTime: <Now>
+          lastTransitionTime: <Now>
+          message: RouteController created a route
+          reason: RouteCreated
+          status: "False"
+          type: NetworkUnavailable
+        daemonEndpoints:
+          kubeletEndpoint:
+            Port: <NodePort>
+        phase: Running
+    kind: patch
+    subresource: status
+    type: application/merge-patch+json
+  stage: node-initialize
diff --git a/kustomize/stage/pod/fast/testdata/pod-pending.input.yaml b/kustomize/stage/pod/fast/testdata/pod-pending.input.yaml
new file mode 100644
index 00000000..f5749c51
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-pending.input.yaml
@@ -0,0 +1,12 @@
+# @Stage: ../pod-ready.yaml
+# @Stage: ../pod-delete.yaml
+# @Stage: ../pod-complete.yaml
+apiVersion: v1
+kind: Pod
+metadata:
+  name: pod-pending
+spec:
+  containers:
+  - name: container
+    image: image
+  nodeName: node
diff --git a/kustomize/stage/pod/fast/testdata/pod-pending.output.yaml b/kustomize/stage/pod/fast/testdata/pod-pending.output.yaml
new file mode 100644
index 00000000..c9ed7107
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-pending.output.yaml
@@ -0,0 +1,34 @@
+apiGroup: v1
+kind: Pod
+name: pod-pending
+stages:
+- next:
+  - data:
+      status:
+        conditions:
+        - lastTransitionTime: <Now>
+          status: "True"
+          type: Initialized
+        - lastTransitionTime: <Now>
+          status: "True"
+          type: Ready
+        - lastTransitionTime: <Now>
+          status: "True"
+          type: ContainersReady
+        containerStatuses:
+        - image: image
+          name: container
+          ready: true
+          restartCount: 0
+          state:
+            running:
+              startedAt: <Now>
+        hostIP: <NodeIPWith("node")>
+        initContainerStatuses: null
+        phase: Running
+        podIP: <PodIPWith("node", false, "", "pod-pending", "")>
+        startTime: <Now>
+    kind: patch
+    subresource: status
+    type: application/merge-patch+json
+  stage: pod-ready
diff --git a/kustomize/stage/pod/fast/testdata/pod-running-with-job.input.yaml b/kustomize/stage/pod/fast/testdata/pod-running-with-job.input.yaml
new file mode 100644
index 00000000..65e6f56a
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-running-with-job.input.yaml
@@ -0,0 +1,28 @@
+# @Stage: ../pod-ready.yaml
+# @Stage: ../pod-delete.yaml
+# @Stage: ../pod-complete.yaml
+apiVersion: v1
+kind: Pod
+metadata:
+  name: pod-running-with-job
+  ownerReferences:
+  - apiVersion: batch/v1
+    kind: Job
+    name: job
+    uid: uid
+spec:
+  containers:
+  - name: container
+    image: image
+  nodeName: node
+status:
+  containerStatuses:
+  - image: image
+    name: container
+    ready: true
+    restartCount: 0
+    state:
+      running:
+        startedAt: <Now>
+  podIP: 10.0.0.1
+  phase: Running
diff --git a/kustomize/stage/pod/fast/testdata/pod-running-with-job.output.yaml b/kustomize/stage/pod/fast/testdata/pod-running-with-job.output.yaml
new file mode 100644
index 00000000..bbca0312
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-running-with-job.output.yaml
@@ -0,0 +1,24 @@
+apiGroup: v1
+kind: Pod
+name: pod-running-with-job
+stages:
+- next:
+  - data:
+      status:
+        containerStatuses:
+        - image: image
+          name: container
+          ready: false
+          restartCount: 0
+          started: false
+          state:
+            terminated:
+              exitCode: 0
+              finishedAt: <Now>
+              reason: Completed
+              startedAt: <Now>
+        phase: Succeeded
+    kind: patch
+    subresource: status
+    type: application/merge-patch+json
+  stage: pod-complete
diff --git a/kustomize/stage/pod/fast/testdata/pod-terminating.input.yaml b/kustomize/stage/pod/fast/testdata/pod-terminating.input.yaml
new file mode 100644
index 00000000..7dce7a5d
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-terminating.input.yaml
@@ -0,0 +1,13 @@
+# @Stage: ../pod-ready.yaml
+# @Stage: ../pod-delete.yaml
+# @Stage: ../pod-complete.yaml
+apiVersion: v1
+kind: Pod
+metadata:
+  name: pod-terminating
+  deletionTimestamp: "2006-01-02T15:04:05Z"
+spec:
+  containers:
+  - name: container
+    image: image
+  nodeName: node
diff --git a/kustomize/stage/pod/fast/testdata/pod-terminating.output.yaml b/kustomize/stage/pod/fast/testdata/pod-terminating.output.yaml
new file mode 100644
index 00000000..1555e9ed
--- /dev/null
+++ b/kustomize/stage/pod/fast/testdata/pod-terminating.output.yaml
@@ -0,0 +1,7 @@
+apiGroup: v1
+kind: Pod
+name: pod-terminating
+stages:
+- next:
+  - kind: delete
+  stage: pod-delete
diff --git a/pkg/kwok/controllers/node_controller_test.go b/pkg/kwok/controllers/node_controller_test.go
index bbf030b6..8871ba88 100644
--- a/pkg/kwok/controllers/node_controller_test.go
+++ b/pkg/kwok/controllers/node_controller_test.go
@@ -84,7 +84,6 @@ func TestNodeController(t *testing.T) {
 		TypedClient:          clientset,
 		NodeIP:               "10.0.0.1",
 		Lifecycle:            resources.NewStaticGetter(lc),
-		FuncMap:              defaultFuncMap,
 		PlayStageParallelism: 2,
 	})
 	if err != nil {
diff --git a/pkg/kwok/controllers/pod_controller_test.go b/pkg/kwok/controllers/pod_controller_test.go
index 8e60111a..7d8e2f89 100644
--- a/pkg/kwok/controllers/pod_controller_test.go
+++ b/pkg/kwok/controllers/pod_controller_test.go
@@ -201,7 +201,6 @@ func TestPodController(t *testing.T) {
 		DisregardStatusWithAnnotationSelector: annotationSelector.String(),
 		Lifecycle:                             resources.NewStaticGetter(lc),
 		NodeGetFunc:                           nodeGetFunc,
-		FuncMap:                               defaultFuncMap,
 		PlayStageParallelism:                  2,
 	})
 	if err != nil {
diff --git a/pkg/kwok/controllers/stage_controller_test.go b/pkg/kwok/controllers/stage_controller_test.go
index 0c38d791..71ed4fa0 100644
--- a/pkg/kwok/controllers/stage_controller_test.go
+++ b/pkg/kwok/controllers/stage_controller_test.go
@@ -88,7 +88,6 @@ func TestStageController(t *testing.T) {
 		DynamicClient:        client,
 		Schema:               patchMeta,
 		Lifecycle:            resources.NewStaticGetter(lc),
-		FuncMap:              defaultFuncMap,
 	})
 	if err != nil {
 		t.Fatal(err)