go,readme: more example & main test code tweaks

anchordotdev · Apr 12, 2024 · 0053ea2 · 0053ea2
1 parent 1998cba
commit 0053ea2
Show file tree

Hide file tree

Showing 3 changed files with 56 additions and 58 deletions.
diff --git a/README.md b/README.md
@@ -30,45 +30,46 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	// wrap os.Exit in a function call so that defer's fire before process exit.
+
+	os.Exit(func() int {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
 
-	// create a test server for each handler to test, this one sets the
-	// response body to the request's Via header
+		// create a test server for each handler to test, this one adds a custom
+		// "Server" header to the response
 
-	srvVia := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Add("Server", "test-via")
-	}))
-	defer srvVia.Close()
+		srvCustom := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Add("Server", "httptest-server")
+		}))
+		defer srvCustom.Close()
 
-	// create fastly config, add a backend for each test server
+		// create fastly config, add a backend for each test server or test case
 
-	cfg := Config{
-		LocalServer: LocalServer{
-			Backends: map[string]Backend{
-				"test-via": {
-					URL: srvVia.URL,
+		cfg := Config{
+			LocalServer: LocalServer{
+				Backends: map[string]Backend{
+					"test-via": {URL: srvCustom.URL},
 				},
 			},
-		},
-	}
+		}
 
-	// create viceroy runner, set the config
+		// create viceroy runner, set the config
 
-	vic, err := NewViceroy(cfg)
-	if err != nil {
-		panic(err)
-	}
-	defer vic.Cleanup()
-
-	// exit with the same code after the tests have run via viceroy to indicate pass/fail
+		vic, err := NewViceroy(cfg)
+		if err != nil {
+			panic(err)
+		}
+		defer vic.Cleanup()
 
-	os.Exit(func() int {
 		// execute the go test command for this package via viceroy
+
 		if err = vic.GoTestPkg(ctx, "fastlytest").Run(); err == nil {
 			return 0
 		}
 
+		// exit with the same code after the tests have run via viceroy to indicate pass/fail
+
 		var eerr *exec.ExitError
 		if errors.Is(err, eerr) {
 			return eerr.ProcessState.ExitCode()
@@ -101,8 +102,6 @@ func TestVia(t *testing.T) {
 	const hdrVia = "1.1 viceroy-test-vm"
 
 	hdlVia := fsthttp.HandlerFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
-		r.Header.Add("Via", hdrVia)
-
 		res, err := r.Send(ctx, "test-via")
 		if err != nil {
 			fsthttp.Error(w, err.Error(), 500)
@@ -134,9 +133,9 @@ func TestVia(t *testing.T) {
 		t.Errorf("want via header %q, got %q", want, got)
 	}
 
-	// assert the header set in srvVia
+	// assert the header set in srvCustom
 
-	if want, got := "test-via", w.HeaderMap.Get("Server"); want != got {
+	if want, got := "httptest-server", w.HeaderMap.Get("Server"); want != got {
 		t.Errorf("want server header %q, got %q", want, got)
 	}
 }

diff --git a/example_test.go b/example_test.go
@@ -15,8 +15,6 @@ func TestVia(t *testing.T) {
 	const hdrVia = "1.1 viceroy-test-vm"
 
 	hdlVia := fsthttp.HandlerFunc(func(ctx context.Context, w fsthttp.ResponseWriter, r *fsthttp.Request) {
-		r.Header.Add("Via", hdrVia)
-
 		res, err := r.Send(ctx, "test-via")
 		if err != nil {
 			fsthttp.Error(w, err.Error(), 500)
@@ -48,9 +46,9 @@ func TestVia(t *testing.T) {
 		t.Errorf("want via header %q, got %q", want, got)
 	}
 
-	// assert the header set in srvVia
+	// assert the header set in srvCustom
 
-	if want, got := "test-via", w.HeaderMap.Get("Server"); want != got {
+	if want, got := "httptest-server", w.HeaderMap.Get("Server"); want != got {
 		t.Errorf("want server header %q, got %q", want, got)
 	}
 }
diff --git a/main_test.go b/main_test.go
@@ -15,45 +15,46 @@ import (
 )
 
 func TestMain(m *testing.M) {
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
+	// wrap os.Exit in a function call so that defer's fire before process exit.
 
-	// create a test server for each handler to test, this one sets the
-	// response body to the request's Via header
+	os.Exit(func() int {
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		// create a test server for each handler to test, this one adds a custom
+		// "Server" header to the response
 
-	srvVia := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Header().Add("Server", "test-via")
-	}))
-	defer srvVia.Close()
+		srvCustom := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Add("Server", "httptest-server")
+		}))
+		defer srvCustom.Close()
 
-	// create fastly config, add a backend for each test server
+		// create fastly config, add a backend for each test server or test case
 
-	cfg := Config{
-		LocalServer: LocalServer{
-			Backends: map[string]Backend{
-				"test-via": {
-					URL: srvVia.URL,
+		cfg := Config{
+			LocalServer: LocalServer{
+				Backends: map[string]Backend{
+					"test-via": {URL: srvCustom.URL},
 				},
 			},
-		},
-	}
-
-	// create viceroy runner, set the config
+		}
 
-	vic, err := NewViceroy(cfg)
-	if err != nil {
-		panic(err)
-	}
-	defer vic.Cleanup()
+		// create viceroy runner, set the config
 
-	// exit with the same code after the tests have run via viceroy to indicate pass/fail
+		vic, err := NewViceroy(cfg)
+		if err != nil {
+			panic(err)
+		}
+		defer vic.Cleanup()
 
-	os.Exit(func() int {
 		// execute the go test command for this package via viceroy
+
 		if err = vic.GoTestPkg(ctx, "fastlytest").Run(); err == nil {
 			return 0
 		}
 
+		// exit with the same code after the tests have run via viceroy to indicate pass/fail
+
 		var eerr *exec.ExitError
 		if errors.Is(err, eerr) {
 			return eerr.ProcessState.ExitCode()