From 31503b121d5ad7971e88a40ebd662c32c0ce01f6 Mon Sep 17 00:00:00 2001
From: Derek McGowan <derek@mcgstyle.net>
Date: Fri, 27 Feb 2015 18:57:34 -0800
Subject: [PATCH] Fix unit test races

Signed-off-by: Derek McGowan <derek@mcgstyle.net> (github: dmcgowan)
Signed-off-by: Sebastiaan van Stijn <github@gone.nl>
---
 spdy_test.go  | 27 ++++++++++++++++++---------
 ws/ws_test.go |  2 +-
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/spdy_test.go b/spdy_test.go
index 312a950..598604e 100644
--- a/spdy_test.go
+++ b/spdy_test.go
@@ -50,7 +50,7 @@ func TestSpdyStreams(t *testing.T) {
 	}
 	go spdyConn.Serve(NoOpStreamHandler)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 	if streamErr != nil {
 		t.Fatalf("Error creating stream: %s", streamErr)
@@ -144,7 +144,7 @@ func TestSpdyStreams(t *testing.T) {
 		t.Fatalf("Error reseting stream: %s", streamResetErr)
 	}
 
-	authenticated = false
+	setAuthenticated(false)
 	badStream, badStreamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 	if badStreamErr != nil {
 		t.Fatalf("Error creating stream: %s", badStreamErr)
@@ -225,7 +225,7 @@ func TestHalfClose(t *testing.T) {
 	}
 	go spdyConn.Serve(NoOpStreamHandler)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 	if streamErr != nil {
 		t.Fatalf("Error creating stream: %s", streamErr)
@@ -311,7 +311,7 @@ func TestUnexpectedRemoteConnectionClosed(t *testing.T) {
 		}
 		go spdyConn.Serve(NoOpStreamHandler)
 
-		authenticated = true
+		setAuthenticated(true)
 		stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 		if streamErr != nil {
 			t.Fatalf("Error creating stream: %s", streamErr)
@@ -427,7 +427,7 @@ func TestIdleShutdownRace(t *testing.T) {
 	}
 	go spdyConn.Serve(NoOpStreamHandler)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, err := spdyConn.CreateStream(http.Header{}, nil, false)
 	if err != nil {
 		t.Fatalf("Error creating stream: %v", err)
@@ -564,7 +564,7 @@ func TestIdleWithData(t *testing.T) {
 
 	spdyConn.SetIdleTimeout(25 * time.Millisecond)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, err := spdyConn.CreateStream(http.Header{}, nil, false)
 	if err != nil {
 		t.Fatalf("Error creating stream: %v", err)
@@ -626,7 +626,7 @@ func TestIdleRace(t *testing.T) {
 
 	spdyConn.SetIdleTimeout(10 * time.Millisecond)
 
-	authenticated = true
+	setAuthenticated(true)
 
 	for i := 0; i < 10; i++ {
 		_, err := spdyConn.CreateStream(http.Header{}, nil, false)
@@ -711,7 +711,7 @@ func TestStreamReset(t *testing.T) {
 	}
 	go spdyConn.Serve(NoOpStreamHandler)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 	if streamErr != nil {
 		t.Fatalf("Error creating stream: %s", streamErr)
@@ -759,7 +759,7 @@ func TestStreamResetWithDataRemaining(t *testing.T) {
 	}
 	go spdyConn.Serve(NoOpStreamHandler)
 
-	authenticated = true
+	setAuthenticated(true)
 	stream, streamErr := spdyConn.CreateStream(http.Header{}, nil, false)
 	if streamErr != nil {
 		t.Fatalf("Error creating stream: %s", streamErr)
@@ -1153,9 +1153,18 @@ func TestStreamReadUnblocksAfterCloseThenReset(t *testing.T) {
 	}
 }
 
+var authLock sync.RWMutex
 var authenticated bool
 
+func setAuthenticated(b bool) {
+	authLock.Lock()
+	defer authLock.Unlock()
+	authenticated = b
+}
+
 func authStreamHandler(stream *Stream) {
+	authLock.RLock()
+	defer authLock.RUnlock()
 	if !authenticated {
 		stream.Refuse()
 		return
diff --git a/ws/ws_test.go b/ws/ws_test.go
index 32d9f0f..4a7f0f5 100644
--- a/ws/ws_test.go
+++ b/ws/ws_test.go
@@ -25,8 +25,8 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/moby/spdystream"
 	"github.com/gorilla/websocket"
+	"github.com/moby/spdystream"
 )
 
 var upgrader = websocket.Upgrader{