From ce5098922ef86c2d6275086614c57ccaa081dc74 Mon Sep 17 00:00:00 2001
From: "jingdi.zhu" <jingdi.zhu@bytedance.com>
Date: Sun, 19 Jan 2025 22:05:31 +0800
Subject: [PATCH] Add exiter to test os.Exit case

---
 runner/engine.go      | 15 ++++++--
 runner/engine_test.go | 85 ++++++++++++++++++++++++++++++++++++++++++-
 runner/exiter.go      | 13 +++++++
 3 files changed, 109 insertions(+), 4 deletions(-)
 create mode 100644 runner/exiter.go

diff --git a/runner/engine.go b/runner/engine.go
index 46c3c88c..ce2fa843 100644
--- a/runner/engine.go
+++ b/runner/engine.go
@@ -17,7 +17,9 @@ import (
 
 // Engine ...
 type Engine struct {
-	config    *Config
+	config *Config
+
+	exiter    exiter
 	proxy     *Proxy
 	logger    *logger
 	watcher   filenotify.FileWatcher
@@ -50,6 +52,7 @@ func NewEngineWithConfig(cfg *Config, debugMode bool) (*Engine, error) {
 	}
 	e := Engine{
 		config:         cfg,
+		exiter:         defaultExiter{},
 		proxy:          NewProxy(&cfg.Proxy),
 		logger:         logger,
 		watcher:        watcher,
@@ -655,8 +658,14 @@ func (e *Engine) stopBin() {
 			close(exitCode)
 		}
 	})
-	if ret := <-exitCode; ret != 0 {
-		os.Exit(ret)
+
+	select {
+	case ret := <-exitCode:
+		if ret != 0 {
+			e.exiter.Exit(ret) // Use exiter instead of direct os.Exit, it's for tests purpose.
+		}
+	case <-time.After(5 * time.Second):
+		e.mainDebug("timed out waiting for process exit")
 	}
 }
 
diff --git a/runner/engine_test.go b/runner/engine_test.go
index 6ffd42fc..2155743b 100644
--- a/runner/engine_test.go
+++ b/runner/engine_test.go
@@ -1044,7 +1044,7 @@ include_file = ["main.sh"]
 
 	t.Logf("start change main.sh")
 	go func() {
-		err := os.WriteFile("main.sh", []byte("#!/bin/sh\nprintf modified > output"), 0o755)
+		err = os.WriteFile("main.sh", []byte("#!/bin/sh\nprintf modified > output"), 0o755)
 		if err != nil {
 			log.Fatalf("Error updating file: %s.", err)
 		}
@@ -1058,3 +1058,86 @@ include_file = ["main.sh"]
 	}
 	assert.Equal(t, []byte("modified"), bytes)
 }
+
+type testExiter struct {
+	t          *testing.T
+	called     bool
+	expectCode int
+}
+
+func (te *testExiter) Exit(code int) {
+	te.called = true
+	if code != te.expectCode {
+		te.t.Fatalf("expected exit code %d, got %d", te.expectCode, code)
+	}
+}
+
+func TestEngineExit(t *testing.T) {
+	tests := []struct {
+		name       string
+		setup      func(*Engine, chan<- int)
+		expectCode int
+		wantCalled bool
+	}{
+		{
+			name: "normal exit - no error",
+			setup: func(_ *Engine, exitCode chan<- int) {
+				go func() {
+					exitCode <- 0
+				}()
+			},
+			expectCode: 0,
+			wantCalled: false,
+		},
+		{
+			name: "error exit - non-zero code",
+			setup: func(_ *Engine, exitCode chan<- int) {
+				go func() {
+					exitCode <- 1
+				}()
+			},
+			expectCode: 1,
+			wantCalled: true,
+		},
+		{
+			name: "process timeout",
+			setup: func(_ *Engine, _ chan<- int) {
+			},
+			expectCode: 0,
+			wantCalled: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			e, err := NewEngine("", true)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			exiter := &testExiter{
+				t:          t,
+				expectCode: tt.expectCode,
+			}
+			e.exiter = exiter
+
+			exitCode := make(chan int)
+
+			if tt.setup != nil {
+				tt.setup(e, exitCode)
+			}
+			select {
+			case ret := <-exitCode:
+				if ret != 0 {
+					e.exiter.Exit(ret)
+				}
+			case <-time.After(1 * time.Millisecond):
+				// timeout case
+			}
+
+			if tt.wantCalled != exiter.called {
+				t.Errorf("Exit() called = %v, want %v", exiter.called, tt.wantCalled)
+			}
+		})
+	}
+}
diff --git a/runner/exiter.go b/runner/exiter.go
new file mode 100644
index 00000000..21cf850d
--- /dev/null
+++ b/runner/exiter.go
@@ -0,0 +1,13 @@
+package runner
+
+import "os"
+
+type exiter interface {
+	Exit(code int)
+}
+
+type defaultExiter struct{}
+
+func (d defaultExiter) Exit(code int) {
+	os.Exit(code)
+}