Use httptest package and properly close servers

http_client_test.go

@@ -75,6 +75,7 @@ func TestHTTPClientSend(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
 
@@ -103,6 +104,7 @@ func TestHTTPClientResponseBuffer(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: false, ResponseBufferSize: 1024})
 
@@ -153,6 +155,7 @@ func TestHTTPClientHTTPSSend(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{})
 
@@ -246,6 +249,7 @@ func TestHTTPClientRedirect(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 1, Debug: false})
 
@@ -277,6 +281,7 @@ func TestHTTPClientRedirectLimit(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 2, Debug: false})
 
@@ -300,6 +305,7 @@ func TestHTTPClientHandleHTTP10(t *testing.T) {
 
 		wg.Done()
 	}))
+	defer server.Close()
 
 	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
 

input_raw_test.go

@@ -22,11 +22,14 @@ func TestRAWInput(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {})
+	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	defer origin.Close()
+	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
+
 
 	var respCounter, reqCounter int64
 
-	input := NewRAWInput(listener.Addr().String(), testRawExpire)
+	input := NewRAWInput(originAddr, testRawExpire)
 	output := NewTestOutput(func(data []byte) {
 		if data[0] == '1' {
 			atomic.AddInt64(&reqCounter, 1)
@@ -42,9 +45,7 @@ func TestRAWInput(t *testing.T) {
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
 
-	address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1)
-
-	client := NewHTTPClient(address, &HTTPClientConfig{})
+	client := NewHTTPClient(origin.URL, &HTTPClientConfig{})
 
 	time.Sleep(time.Millisecond)
 
@@ -68,14 +69,15 @@ func TestInputRAW100Expect(t *testing.T) {
 	fileContent, _ := ioutil.ReadFile("README.md")
 
 	// Origing and Replay server initialization
-	origin := startHTTP(func(w http.ResponseWriter, req *http.Request) {
-		defer req.Body.Close()
-		ioutil.ReadAll(req.Body)
+	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		ioutil.ReadAll(r.Body)
 
 		wg.Done()
-	})
+	}))
+	defer origin.Close()
 
-	originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1)
+	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 
 	input := NewRAWInput(originAddr, testRawExpire)
 
@@ -92,20 +94,20 @@ func TestInputRAW100Expect(t *testing.T) {
 		}
 	})
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
-		defer req.Body.Close()
-		body, _ := ioutil.ReadAll(req.Body)
+	replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		body, _ := ioutil.ReadAll(r.Body)
 
 		if !bytes.Equal(body, fileContent) {
-			buf, _ := httputil.DumpRequest(req, true)
+			buf, _ := httputil.DumpRequest(r, true)
 			t.Error("Wrong POST body:", string(buf))
 		}
 
 		wg.Done()
-	})
-	replayAddr := listener.Addr().String()
+	}))
+	defer replay.Close()
 
-	httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{})
+	httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{testOutput, httpOutput}
@@ -131,31 +133,30 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 	fileContent, _ := ioutil.ReadFile("README.md")
 
 	// Origing and Replay server initialization
-	origin := startHTTP(func(w http.ResponseWriter, req *http.Request) {
-		defer req.Body.Close()
-		ioutil.ReadAll(req.Body)
+	origin := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		ioutil.ReadAll(r.Body)
 
 		wg.Done()
-	})
-
-	originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1)
+	}))
 
+	originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 	input := NewRAWInput(originAddr, testRawExpire)
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
-		defer req.Body.Close()
-		body, _ := ioutil.ReadAll(req.Body)
+	replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		defer r.Body.Close()
+		body, _ := ioutil.ReadAll(r.Body)
 
 		if !bytes.Equal(body, fileContent) {
-			buf, _ := httputil.DumpRequest(req, true)
+			buf, _ := httputil.DumpRequest(r, true)
 			t.Error("Wrong POST body:", string(buf))
 		}
 
 		wg.Done()
-	})
-	replayAddr := listener.Addr().String()
+	}))
+	defer replay.Close()
 
-	httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{Debug: true})
+	httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{Debug: true})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{httpOutput}

middleware_test.go

@@ -17,7 +17,7 @@ import (
 type fakeServiceCb func(string, int, []byte)
 
 // Simple service that generate token on request, and require this token for accesing to secure area
-func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) string {
+func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) *httptest.Server {
 	active_tokens := make([]string, 0)
 	var mu sync.Mutex
 
@@ -61,8 +61,7 @@ func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) string {
 		wg.Done()
 	}))
 
-	address := strings.Replace(server.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
-	return address
+	return server
 }
 
 func TestFakeSecureService(t *testing.T) {
@@ -70,12 +69,13 @@ func TestFakeSecureService(t *testing.T) {
 
 	wg := new(sync.WaitGroup)
 
-	addr := NewFakeSecureService(wg, func(path string, status int, resp []byte) {
+	server := NewFakeSecureService(wg, func(path string, status int, resp []byte) {
 	})
+	defer server.Close()
 
 	wg.Add(3)
 
-	client := NewHTTPClient("http://"+addr, &HTTPClientConfig{Debug: true})
+	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
 	resp, _ = client.Get("/token")
 	token = proto.Body(resp)
 
@@ -102,11 +102,14 @@ func TestEchoMiddleware(t *testing.T) {
 		w.Header().Set("RequestPath", r.URL.Path)
 		wg.Done()
 	}))
+	defer from.Close()
+
 	to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Env", "test")
 		w.Header().Set("RequestPath", r.URL.Path)
 		wg.Done()
 	}))
+	defer to.Close()
 
 	quit := make(chan int)
 
@@ -150,6 +153,8 @@ func TestTokenMiddleware(t *testing.T) {
 	from := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
 		time.Sleep(10 * time.Millisecond)
 	})
+	defer from.Close()
+
 	to := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
 		switch path {
 		case "/secure":
@@ -160,16 +165,18 @@ func TestTokenMiddleware(t *testing.T) {
 
 		time.Sleep(10 * time.Millisecond)
 	})
+	defer to.Close()
 
 	quit := make(chan int)
 
 	Settings.middleware = "go run ./examples/middleware/token_modifier.go"
 
+	fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
 	// Catch traffic from one service
-	input := NewRAWInput(from, testRawExpire)
+	input := NewRAWInput(fromAddr, testRawExpire)
 
 	// And redirect to another
-	output := NewHTTPOutput(to, &HTTPOutputConfig{Debug: true})
+	output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -184,7 +191,7 @@ func TestTokenMiddleware(t *testing.T) {
 	// Should receive 2 requests from original + 2 from replayed
 	wg.Add(4)
 
-	client := NewHTTPClient("http://"+from, &HTTPClientConfig{Debug: false})
+	client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: false})
 
 	// Sending traffic to original service
 	resp, _ = client.Get("/token")

output_http_test.go

@@ -3,7 +3,6 @@ package main
 import (
 	"io"
 	"io/ioutil"
-	"net"
 	"net/http"
 	"net/http/httptest"
 	_ "net/http/httputil"
@@ -12,25 +11,13 @@ import (
 	"time"
 )
 
-func startHTTP(cb func(http.ResponseWriter, *http.Request)) net.Listener {
-	handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		cb(w, r)
-	})
-
-	listener, _ := net.Listen("tcp", ":0")
-
-	go http.Serve(listener, handler)
-
-	return listener
-}
-
 func TestHTTPOutput(t *testing.T) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
 	input := NewTestInput()
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
+	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if req.Header.Get("User-Agent") != "Gor" {
 			t.Error("Wrong header")
 		}
@@ -49,13 +36,14 @@ func TestHTTPOutput(t *testing.T) {
 		}
 
 		wg.Done()
-	})
+	}))
+	defer server.Close()
 
 	headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}}
 	methods := HTTPMethods{[]byte("GET"), []byte("PUT"), []byte("POST")}
 	Settings.modifierConfig = HTTPModifierConfig{headers: headers, methods: methods}
 
-	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{Debug: true})
+	output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: true})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -82,18 +70,19 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
 
 	input := NewTestInput()
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
+	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if req.Host != "custom-host.com" {
 			t.Error("Wrong header", req.Host)
 		}
 
 		wg.Done()
-	})
+	}))
+	defer server.Close()
 
 	headers := HTTPHeaders{HTTPHeader{"Host", "custom-host.com"}}
 	Settings.modifierConfig = HTTPModifierConfig{headers: headers}
 
-	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{Debug: false, OriginalHost: true})
+	output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: false, OriginalHost: true})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}
@@ -140,13 +129,14 @@ func BenchmarkHTTPOutput(b *testing.B) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
-	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
+	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		time.Sleep(50 * time.Millisecond)
 		wg.Done()
-	})
+	}))
+	defer server.Close()
 
 	input := NewTestInput()
-	output := NewHTTPOutput(listener.Addr().String(), &HTTPOutputConfig{})
+	output := NewHTTPOutput(server.URL, &HTTPOutputConfig{})
 
 	Plugins.Inputs = []io.Reader{input}
 	Plugins.Outputs = []io.Writer{output}