From b0e39794371d49895e3fc8db8307d6971e825d19 Mon Sep 17 00:00:00 2001
From: kiyon <kiyonlin@163.com>
Date: Wed, 15 Jul 2020 10:37:57 +0800
Subject: [PATCH] =?UTF-8?q?=F0=9F=91=8D=20improve=20prefork=20test=20cover?=
 =?UTF-8?q?age?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 prefork.go      | 21 ++++++++++++++-------
 prefork_test.go | 42 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 7 deletions(-)
 create mode 100644 prefork_test.go

diff --git a/prefork.go b/prefork.go
index a8e2cc9f..67341888 100644
--- a/prefork.go
+++ b/prefork.go
@@ -17,6 +17,8 @@ const (
 	envPreforkChildVal = "1"
 )
 
+var testPreforkMaster = false
+
 // IsChild determines if the current process is a result of Prefork
 func (app *App) IsChild() bool {
 	return os.Getenv(envPreforkChildKey) == envPreforkChildVal
@@ -43,11 +45,12 @@ func (app *App) prefork(addr string, tlsconfig ...*tls.Config) (err error) {
 
 		// kill child proc when master exits
 		go func() {
-			ppid, err := os.FindProcess(os.Getppid())
+			p, err := os.FindProcess(os.Getppid())
 			if err == nil {
-				_, _ = ppid.Wait()
+				_, _ = p.Wait()
+			} else {
+				os.Exit(1)
 			}
-			os.Exit(1)
 		}()
 		// listen for incoming connections
 		return app.server.Serve(ln)
@@ -71,12 +74,15 @@ func (app *App) prefork(addr string, tlsconfig ...*tls.Config) (err error) {
 	}()
 
 	// collect child pids
-	pids := []string{}
+	var pids []string
 
 	// launch child procs
 	for i := 0; i < max; i++ {
 		/* #nosec G204 */
 		cmd := exec.Command(os.Args[0], os.Args[1:]...)
+		if testPreforkMaster {
+			cmd = exec.Command("cd")
+		}
 		cmd.Stdout = os.Stdout
 		cmd.Stderr = os.Stderr
 
@@ -89,12 +95,13 @@ func (app *App) prefork(addr string, tlsconfig ...*tls.Config) (err error) {
 		}
 
 		// store child process
-		childs[cmd.Process.Pid] = cmd
-		pids = append(pids, strconv.Itoa(cmd.Process.Pid))
+		pid := cmd.Process.Pid
+		childs[pid] = cmd
+		pids = append(pids, strconv.Itoa(pid))
 
 		// notify master if child crashes
 		go func() {
-			channel <- child{cmd.Process.Pid, cmd.Wait()}
+			channel <- child{pid, cmd.Wait()}
 		}()
 	}
 
diff --git a/prefork_test.go b/prefork_test.go
new file mode 100644
index 00000000..d06117a2
--- /dev/null
+++ b/prefork_test.go
@@ -0,0 +1,42 @@
+package fiber
+
+import (
+	"os"
+	"testing"
+	"time"
+
+	utils "github.com/gofiber/utils"
+)
+
+func Test_App_Prefork_Child_Process(t *testing.T) {
+	utils.AssertEqual(t, nil, os.Setenv(envPreforkChildKey, envPreforkChildVal))
+	defer os.Setenv(envPreforkChildKey, "")
+
+	app := New(&Settings{
+		DisableStartupMessage: true,
+	})
+	app.init()
+
+	go func() {
+		time.Sleep(1000 * time.Millisecond)
+		utils.AssertEqual(t, nil, app.Shutdown())
+	}()
+
+	utils.AssertEqual(t, nil, app.prefork("127.0.0.1:"))
+}
+
+func Test_App_Prefork_Main_Process(t *testing.T) {
+	testPreforkMaster = true
+
+	app := New(&Settings{
+		DisableStartupMessage: true,
+	})
+	app.init()
+
+	go func() {
+		time.Sleep(1000 * time.Millisecond)
+		utils.AssertEqual(t, nil, app.Shutdown())
+	}()
+
+	utils.AssertEqual(t, nil, app.prefork("127.0.0.1:"))
+}