Support only 1 middleware and rewrite tests

middleware_test.go

@@ -1,22 +1,20 @@
 package main
 
 import (
-	_ "bufio"
 	"bytes"
 	"crypto/rand"
 	"io"
-	"io/ioutil"
-	_ "log"
-	_ "net"
 	"net/http"
 	"sync"
 	"testing"
     "strings"
+    "github.com/buger/gor/proto"
+    "encoding/hex"
 )
 
 // Simple service that generate token on request, and require this token for accesing to secure area
 func NewFakeSecureService(wg *sync.WaitGroup) string {
-	active_tokens := make([][]byte, 0)
+	active_tokens := make([]string, 0)
 
 	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
         Debug("Received request: " + req.URL.String())
@@ -25,17 +23,18 @@ func NewFakeSecureService(wg *sync.WaitGroup) string {
 		case "/token":
 			// Generate random token
 			token_length := 10
-			token := make([]byte, token_length)
-			rand.Read(token)
+			buf := make([]byte, token_length)
+			rand.Read(buf)
+            token := hex.EncodeToString(buf)
 			active_tokens = append(active_tokens, token)
 
-            w.Write(token)
+            w.Write([]byte(token))
 		case "/secure":
-			token := []byte(req.URL.Query().Get("token"))
+			token := req.URL.Query().Get("token")
             token_found := false
 
 			for _, t := range active_tokens {
-				if bytes.Equal(t, token) {
+				if t == token {
                     token_found = true
                     break
 				}
@@ -56,7 +55,7 @@ func NewFakeSecureService(wg *sync.WaitGroup) string {
 }
 
 func TestFakeSecureService(t *testing.T) {
-	var resp *http.Response
+    var resp, token []byte
 
 	wg := new(sync.WaitGroup)
 
@@ -64,26 +63,27 @@ func TestFakeSecureService(t *testing.T) {
 
 	wg.Add(3)
 
-	resp, _ = http.Get("http://" + addr + "/token")
-	token, _ := ioutil.ReadAll(resp.Body)
+    client := NewHTTPClient("http://" + addr, &HTTPClientConfig{Debug: true})
+    resp, _ = client.Get("/token")
+    token = proto.Body(resp)
 
-	// Right token
-	resp, _ = http.Get("http://" + addr + "/secure?token=" + string(token))
-	if resp.StatusCode != http.StatusAccepted {
-		t.Error("Valid token should returns wrong status:", resp.StatusCode)
-	}
+    // Right token
+    resp, _ = client.Get("/secure?token=" + string(token))
+    if !bytes.Equal(proto.Status(resp), []byte("202")) {
+        t.Error("Valid token should return status 202:", string(proto.Status(resp)))
+    }
 
-	// Wrong tokens forbidden
-	resp, _ = http.Get("http://" + addr + "/secure?token=wrong")
-	if resp.StatusCode != http.StatusForbidden {
-		t.Error("Wrong tokens should be forbidden, instead:", resp.StatusCode)
+    // Wrong tokens forbidden
+    resp, _ = client.Get("/secure?token=wrong")
+    if !bytes.Equal(proto.Status(resp), []byte("403")) {
+		t.Error("Wrong token should returns status 403:", string(proto.Status(resp)))
 	}
 
 	wg.Wait()
 }
 
 func TestMiddleware(t *testing.T) {
-    var resp *http.Response
+    var resp, token []byte
 
     wg := new(sync.WaitGroup)
 
@@ -107,13 +107,15 @@ func TestMiddleware(t *testing.T) {
     // Should receive 2 requests from original + 2 from replayed
     wg.Add(4)
 
-    // Sending traffic to original service
-    resp, _ = http.Get("http://" + from + "/token")
-    token, _ := ioutil.ReadAll(resp.Body)
+    client := NewHTTPClient("http://" + from, &HTTPClientConfig{Debug: true})
 
-    resp, _ = http.Get("http://" + from + "/secure?token=" + string(token))
-    if resp.StatusCode != http.StatusAccepted {
-        t.Error("Valid token should returns wrong status:", resp.StatusCode)
+    // Sending traffic to original service
+    resp, _ = client.Get("/token")
+    token = proto.Body(resp)
+
+    resp, _ = client.Get("/secure?token=" + string(token))
+    if !bytes.Equal(proto.Status(resp), []byte("202")) {
+        t.Error("Valid token should return 202:", proto.Status(resp))
     }
 
     wg.Wait()