From 513100bb9227390ace03ce48767e61cb43ec9549 Mon Sep 17 00:00:00 2001
From: Jakob Borg <jakob@nym.se>
Date: Sat, 22 Mar 2014 17:06:15 +0100
Subject: [PATCH] Fix tests for >1 CPU (fixes #99)

---
 build.sh                  |  2 +-
 protocol/common_test.go   | 34 ++++++++++++++++++++++++++--------
 protocol/protocol_test.go | 31 +++++++++++++++----------------
 3 files changed, 42 insertions(+), 25 deletions(-)

diff --git a/build.sh b/build.sh
index a2b945d8..af31c4ac 100755
--- a/build.sh
+++ b/build.sh
@@ -15,7 +15,7 @@ prepare() {
 }
 
 test() {
-	go test ./...
+	go test -cpu=1,2,4 ./...
 }
 
 tarDist() {
diff --git a/protocol/common_test.go b/protocol/common_test.go
index 469c271e..a24075e0 100644
--- a/protocol/common_test.go
+++ b/protocol/common_test.go
@@ -1,14 +1,23 @@
 package protocol
 
-import "io"
+import (
+	"io"
+	"time"
+)
 
 type TestModel struct {
-	data   []byte
-	repo   string
-	name   string
-	offset int64
-	size   int
-	closed bool
+	data     []byte
+	repo     string
+	name     string
+	offset   int64
+	size     int
+	closedCh chan bool
+}
+
+func newTestModel() *TestModel {
+	return &TestModel{
+		closedCh: make(chan bool),
+	}
 }
 
 func (t *TestModel) Index(nodeID string, files []FileInfo) {
@@ -26,7 +35,16 @@ func (t *TestModel) Request(nodeID, repo, name string, offset int64, size int) (
 }
 
 func (t *TestModel) Close(nodeID string, err error) {
-	t.closed = true
+	close(t.closedCh)
+}
+
+func (t *TestModel) isClosed() bool {
+	select {
+	case <-t.closedCh:
+		return true
+	case <-time.After(1 * time.Second):
+		return false // Timeout
+	}
 }
 
 type ErrPipe struct {
diff --git a/protocol/protocol_test.go b/protocol/protocol_test.go
index c8f9cd3f..50a968ac 100644
--- a/protocol/protocol_test.go
+++ b/protocol/protocol_test.go
@@ -5,7 +5,6 @@ import (
 	"io"
 	"testing"
 	"testing/quick"
-	"time"
 )
 
 func TestHeaderFunctions(t *testing.T) {
@@ -42,8 +41,8 @@ func TestPingErr(t *testing.T) {
 
 	for i := 0; i < 12; i++ {
 		for j := 0; j < 12; j++ {
-			m0 := &TestModel{}
-			m1 := &TestModel{}
+			m0 := newTestModel()
+			m1 := newTestModel()
 
 			ar, aw := io.Pipe()
 			br, bw := io.Pipe()
@@ -69,8 +68,9 @@ func TestRequestResponseErr(t *testing.T) {
 	var pass bool
 	for i := 0; i < 48; i++ {
 		for j := 0; j < 38; j++ {
-			m0 := &TestModel{data: []byte("response data")}
-			m1 := &TestModel{}
+			m0 := newTestModel()
+			m0.data = []byte("response data")
+			m1 := newTestModel()
 
 			ar, aw := io.Pipe()
 			br, bw := io.Pipe()
@@ -83,11 +83,10 @@ func TestRequestResponseErr(t *testing.T) {
 			d, err := c1.Request("default", "tn", 1234, 5678)
 			if err == e || err == ErrClosed {
 				t.Logf("Error at %d+%d bytes", i, j)
-				if !m1.closed {
+				if !m1.isClosed() {
 					t.Error("c1 not closed")
 				}
-				time.Sleep(1 * time.Millisecond)
-				if !m0.closed {
+				if !m0.isClosed() {
 					t.Error("c0 not closed")
 				}
 				continue
@@ -120,8 +119,8 @@ func TestRequestResponseErr(t *testing.T) {
 }
 
 func TestVersionErr(t *testing.T) {
-	m0 := &TestModel{}
-	m1 := &TestModel{}
+	m0 := newTestModel()
+	m1 := newTestModel()
 
 	ar, aw := io.Pipe()
 	br, bw := io.Pipe()
@@ -136,14 +135,14 @@ func TestVersionErr(t *testing.T) {
 	}))
 	c0.flush()
 
-	if !m1.closed {
+	if !m1.isClosed() {
 		t.Error("Connection should close due to unknown version")
 	}
 }
 
 func TestTypeErr(t *testing.T) {
-	m0 := &TestModel{}
-	m1 := &TestModel{}
+	m0 := newTestModel()
+	m1 := newTestModel()
 
 	ar, aw := io.Pipe()
 	br, bw := io.Pipe()
@@ -158,14 +157,14 @@ func TestTypeErr(t *testing.T) {
 	}))
 	c0.flush()
 
-	if !m1.closed {
+	if !m1.isClosed() {
 		t.Error("Connection should close due to unknown message type")
 	}
 }
 
 func TestClose(t *testing.T) {
-	m0 := &TestModel{}
-	m1 := &TestModel{}
+	m0 := newTestModel()
+	m1 := newTestModel()
 
 	ar, aw := io.Pipe()
 	br, bw := io.Pipe()