From 9dd6a107d8fc9c9a308961d4e316c6a0211c392a Mon Sep 17 00:00:00 2001
From: iamqizhao <toqizhao@gmail.com>
Date: Fri, 3 Apr 2015 17:53:58 -0700
Subject: [PATCH] fix a race in transport_test.go

---
 transport/transport_test.go | 34 +++++++++++++++++++++++++---------
 1 file changed, 25 insertions(+), 9 deletions(-)

diff --git a/transport/transport_test.go b/transport/transport_test.go
index ecd1b204..ea6fa698 100644
--- a/transport/transport_test.go
+++ b/transport/transport_test.go
@@ -430,12 +430,22 @@ func TestServerWithMisbehavedClient(t *testing.T) {
 		Method: "foo",
 	}
 	var sc *http2Server
-	for k, _ := range server.conns {
-		var ok bool
-		sc, ok = k.(*http2Server)
-		if !ok {
-			t.Fatalf("Failed to convert %v to *http2Server", k)
+	// Wait until the server transport is setup.
+	for {
+		server.mu.Lock()
+		if len(server.conns) > 0 {
+			for k, _ := range server.conns {
+				var ok bool
+				sc, ok = k.(*http2Server)
+				if !ok {
+					t.Fatalf("Failed to convert %v to *http2Server", k)
+				}
+			}
+			server.mu.Unlock()
+			break
 		}
+		server.mu.Unlock()
+		time.Sleep(time.Millisecond)
 	}
 	cc, ok := ct.(*http2Client)
 	if !ok {
@@ -454,17 +464,23 @@ func TestServerWithMisbehavedClient(t *testing.T) {
 	}
 	cc.writableChan <- 0
 	sent += http2MaxFrameLen
-	// Wait until the server creates the corresponding stream.
+	// Wait until the server creates the corresponding stream and receive some data.
 	var ss *Stream
 	for {
 		time.Sleep(time.Millisecond)
 		sc.mu.Lock()
-		if len(sc.activeStreams) > 0 {
-			ss = sc.activeStreams[s.id]
+		if len(sc.activeStreams) == 0 {
 			sc.mu.Unlock()
+			continue
+		}
+		ss = sc.activeStreams[s.id]
+		sc.mu.Unlock()
+		ss.fc.mu.Lock()
+		if ss.fc.pendingData > 0 {
+			ss.fc.mu.Unlock()
 			break
 		}
-		sc.mu.Unlock()
+		ss.fc.mu.Unlock()
 	}
 	if ss.fc.pendingData != http2MaxFrameLen || ss.fc.pendingUpdate != 0 || sc.fc.pendingData != http2MaxFrameLen || sc.fc.pendingUpdate != 0 {
 		t.Fatalf("Server mistakenly updates inbound flow control params: got %d, %d, %d, %d; want %d, %d, %d, %d", ss.fc.pendingData, ss.fc.pendingUpdate, sc.fc.pendingData, sc.fc.pendingUpdate, http2MaxFrameLen, 0, http2MaxFrameLen, 0)