From 70b0316dcdbe482ec16e81615df3491da0e0e11b Mon Sep 17 00:00:00 2001
From: "F. Gold" <fgold@vmware.com>
Date: Fri, 8 Dec 2023 12:20:29 -0800
Subject: [PATCH] Gracefully exit when setup fails

When an error happens during a setup step, instead of calling klog.Fatal
which causes os.Exit, instead do klog.Error, break, run defer function
with clean up steps, and return.

Include setup fail test example in examples directory.

Signed-off-by: F. Gold <fgold@vmware.com>
---
 examples/graceful_setup_fail/main_test.go | 126 ++++++++++++++++++++++
 pkg/env/env.go                            |   3 +-
 2 files changed, 128 insertions(+), 1 deletion(-)
 create mode 100644 examples/graceful_setup_fail/main_test.go

diff --git a/examples/graceful_setup_fail/main_test.go b/examples/graceful_setup_fail/main_test.go
new file mode 100644
index 00000000..2572729f
--- /dev/null
+++ b/examples/graceful_setup_fail/main_test.go
@@ -0,0 +1,126 @@
+/*
+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 setupfail
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"testing"
+
+	"k8s.io/klog/v2"
+	"sigs.k8s.io/e2e-framework/pkg/env"
+	"sigs.k8s.io/e2e-framework/pkg/envconf"
+)
+
+func TestMain(m *testing.M) {
+	var testEnv env.Environment
+
+	tests := []struct {
+		name             string
+		testEnvGenerator func(context.Context, *[]string) env.Environment
+		expected         []string
+	}{
+		{
+			name: "No test setup failures, do all setup and finish actions",
+			testEnvGenerator: func(ctx context.Context, actions *[]string) env.Environment {
+				testEnv = env.New()
+
+				testEnv.Setup(
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 1")
+						return ctx, nil
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 2")
+						return ctx, nil
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 3")
+						return ctx, nil
+					},
+				)
+				testEnv.Finish(
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed finish 1")
+						return ctx, nil
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed finish 2")
+						return ctx, nil
+					},
+				)
+
+				return testEnv
+			},
+			expected: []string{"completed setup 1", "completed setup 2", "completed setup 3", "completed finish 1", "completed finish 2"},
+		},
+		{
+			name: "Skip remaining setup actions after one fails, do finish actions",
+			testEnvGenerator: func(ctx context.Context, actions *[]string) env.Environment {
+				testEnv = env.New()
+
+				testEnv.Setup(
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 1")
+						return ctx, nil
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 2")
+						return ctx, fmt.Errorf("failed setup 2")
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed setup 3")
+						return ctx, nil
+					},
+				)
+				testEnv.Finish(
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed finish 1")
+						return ctx, nil
+					},
+					func(ctx context.Context, _ *envconf.Config) (context.Context, error) {
+						*actions = append(*actions, "completed finish 2")
+						return ctx, nil
+					},
+				)
+
+				return testEnv
+			},
+			expected: []string{"completed setup 1", "completed setup 2", "completed finish 1", "completed finish 2"},
+		},
+	}
+
+	for _, test := range tests {
+		var actions []string
+		testEnv = test.testEnvGenerator(context.TODO(), &actions)
+		_ = testEnv.Run(m)
+
+		readableActions := strings.Join(actions, ", ")
+		if len(actions) != len(test.expected) {
+			klog.Fatalf("Expected slice of %v, but got %s", test.expected, readableActions)
+		}
+
+		for i := range actions {
+			if actions[i] != test.expected[i] {
+				klog.Fatalf("Expected slice of %v, but got %s", test.expected, readableActions)
+			}
+		}
+
+		klog.Infof("PASS: %s\n\tActions completed: %+v", test.name, readableActions)
+	}
+}
diff --git a/pkg/env/env.go b/pkg/env/env.go
index 3b1e4838..48dd3aef 100644
--- a/pkg/env/env.go
+++ b/pkg/env/env.go
@@ -381,7 +381,8 @@ func (e *testEnv) Run(m *testing.M) (exitCode int) {
 	for _, setup := range setups {
 		// context passed down to each setup
 		if ctx, err = setup.run(ctx, e.cfg); err != nil {
-			klog.Fatalf("%s failure: %s", setup.role, err)
+			klog.Errorf("%s failure: %s", setup.role, err)
+			break
 		}
 	}
 	e.ctx = ctx