MockNetConn: Wrap all channel reads in selects; send EOF on Close(); kill goroutines correctly.

client/mocknetconn_test.go

@@ -4,6 +4,7 @@ import (
 	"net"
 	"os"
 	"testing"
+	"time"
 )
 
 type mockNetConn struct {
@@ -11,6 +12,8 @@ type mockNetConn struct {
 
 	In, Out chan string
 	in, out chan []byte
+	closers []chan bool
+	rc chan bool
 
 	closed bool
 	rt, wt int64
@@ -19,6 +22,7 @@ type mockNetConn struct {
 func MockNetConn(t *testing.T) (*mockNetConn) {
 	// Our mock connection is a testing object
 	m := &mockNetConn{T: t}
+	m.closers = make([]chan bool, 0, 3)
 
 	// set known values for conn info
 	m.closed = false
@@ -28,20 +32,39 @@ func MockNetConn(t *testing.T) (*mockNetConn) {
 	// buffer input
 	m.In = make(chan string, 20)
 	m.in = make(chan []byte)
+	ic := make(chan bool)
+	m.closers = append(m.closers, ic)
 	go func() {
-		for !m.closed {
-			m.in <- []byte(<-m.In)
+		for {
+			select {
+			case <-ic:
+				return
+			case s := <-m.In:
+				m.in <- []byte(s)
+			}
 		}
 	}()
 
 	// buffer output
 	m.Out = make(chan string)
 	m.out = make(chan []byte, 20)
+	oc := make(chan bool)
+	m.closers = append(m.closers, oc)
 	go func() {
-		for !m.closed {
-			m.Out <- string(<-m.out)
+		for {
+			select {
+			case <-oc:
+				return
+			case b := <-m.out:
+				m.Out <- string(b)
+			}
 		}
 	}()
+
+	// Set up channel to force EOF to Read() on close.
+	m.rc = make(chan bool)
+	m.closers = append(m.closers, m.rc)
+
 	return m
 }
 
@@ -51,10 +74,17 @@ func (m *mockNetConn) Send(s string) {
 }
 
 func (m *mockNetConn) Expect(e string) {
-	s := <-m.Out
+	t := time.NewTimer(5e6)
+	select {
+	case <-t.C:
+		m.Errorf("Mock connection did not receive expected output.\n\t" +
+			"Expected: '%s', got nothing.", e)
+	case s := <-m.Out:
+		t.Stop()
 		if e + "\r\n" != s {
 			m.Errorf("Mock connection received unexpected value.\n\t" +
-			"Expected: %s\n\tGot: %s", e, s)
+				"Expected: '%s'\n\tGot: '%s'", e, s)
+		}
 	}
 }
 
@@ -63,9 +93,13 @@ func (m *mockNetConn) Read(b []byte) (int, os.Error) {
 	if m.closed {
 		return 0, os.NewError("EOF")
 	}
-	s := <-m.in
+	select {
+	case s := <-m.in:
 		copy(b, s)
-	return len(s), nil
+	case <-m.rc:
+		return 0, os.EOF
+	}
+	return len(b), nil
 }
 
 func (m *mockNetConn) Write(s []byte) (int, os.Error) {
@@ -79,6 +113,14 @@ func (m *mockNetConn) Write(s []byte) (int, os.Error) {
 }
 
 func (m *mockNetConn) Close() os.Error {
+	if m.closed {
+		return os.EINVAL
+	}
+	// Shut down *ALL* the goroutines!
+	// This will trigger an EOF event in Read() too
+	for _, c := range m.closers {
+		c <- true
+	}
 	m.closed = true
 	return nil
 }