From 5a411fa38e49e5d79602f01fa9aaf058c12e5627 Mon Sep 17 00:00:00 2001
From: Louis Opter <louis@dotcloud.com>
Date: Wed, 10 Jul 2013 18:02:41 -0700
Subject: [PATCH] Make the TestAllocate{UDP,TCP}PortLocalhost more reliable

- For the TCP test try again if socat wasn't listening yet;
- For the UDP test raise the timeout to a minute to workaround what
  seems to be an issue with Linux.
---
 runtime_test.go | 69 ++++++++++++++++++++++++++++++++-----------------
 1 file changed, 46 insertions(+), 23 deletions(-)

diff --git a/runtime_test.go b/runtime_test.go
index 6e64d9e39c..d8033467c2 100644
--- a/runtime_test.go
+++ b/runtime_test.go
@@ -383,31 +383,50 @@ func TestAllocateTCPPortLocalhost(t *testing.T) {
 	defer nuke(runtime)
 	defer container.Kill()
 
-	conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%v", port))
-	if err != nil {
-		t.Fatal(err)
-	}
-	defer conn.Close()
+	for i := 0; i != 10; i++ {
+		conn, err := net.Dial("tcp", fmt.Sprintf("localhost:%v", port))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer conn.Close()
 
-	input := bytes.NewBufferString("well hello there\n")
-	_, err = conn.Write(input.Bytes())
-	if err != nil {
-		t.Fatal(err)
-	}
-	buf := make([]byte, 16)
-	read := 0
-	conn.SetReadDeadline(time.Now().Add(4 * time.Second))
-	read, err = conn.Read(buf)
-	if err != nil {
-		t.Fatal(err)
-	}
-	output := string(buf[:read])
-	if !strings.Contains(output, "well hello there") {
-		t.Fatal(fmt.Errorf("[%v] doesn't contain [well hello there]", output))
+		input := bytes.NewBufferString("well hello there\n")
+		_, err = conn.Write(input.Bytes())
+		if err != nil {
+			t.Fatal(err)
+		}
+		buf := make([]byte, 16)
+		read := 0
+		conn.SetReadDeadline(time.Now().Add(3 * time.Second))
+		read, err = conn.Read(buf)
+		if err != nil {
+			if err, ok := err.(*net.OpError); ok {
+				if err.Err == syscall.ECONNRESET {
+					t.Logf("Connection reset by the proxy, socat is probably not listening yet, trying again in a sec")
+					conn.Close()
+					time.Sleep(time.Second)
+					continue
+				}
+				if err.Timeout() {
+					t.Log("Timeout, trying again")
+					conn.Close()
+					continue
+				}
+			}
+			t.Fatal(err)
+		}
+		output := string(buf[:read])
+		if !strings.Contains(output, "well hello there") {
+			t.Fatal(fmt.Errorf("[%v] doesn't contain [well hello there]", output))
+		} else {
+			return
+		}
 	}
+
+	t.Fatal("No reply from the container")
 }
 
-// Run a container with a TCP port allocated, and test that it can receive connections on localhost
+// Run a container with an UDP port allocated, and test that it can receive connections on localhost
 func TestAllocateUDPPortLocalhost(t *testing.T) {
 	runtime, container, port := startEchoServerContainer(t, "udp")
 	defer nuke(runtime)
@@ -421,12 +440,16 @@ func TestAllocateUDPPortLocalhost(t *testing.T) {
 
 	input := bytes.NewBufferString("well hello there\n")
 	buf := make([]byte, 16)
-	for i := 0; i != 20; i++ {
+	// Try for a minute, for some reason the select in socat may take ages
+	// to return even though everything on the path seems fine (i.e: the
+	// UDPProxy forwards the traffic correctly and you can see the packets
+	// on the interface from within the container).
+	for i := 0; i != 120; i++ {
 		_, err := conn.Write(input.Bytes())
 		if err != nil {
 			t.Fatal(err)
 		}
-		conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+		conn.SetReadDeadline(time.Now().Add(500 * time.Millisecond))
 		read, err := conn.Read(buf)
 		if err == nil {
 			output := string(buf[:read])