From a22d9380e413a23341103a13385227aef729a8ea Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Fri, 9 Sep 2016 12:11:27 -0400
Subject: [PATCH] refactor: fixed failing tests

---
 chat/room.go |  2 +-
 host.go      |  3 +++
 host_test.go | 24 ++++++++++++++++++------
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/chat/room.go b/chat/room.go
index 1bf85db..e94f454 100644
--- a/chat/room.go
+++ b/chat/room.go
@@ -143,7 +143,7 @@ func (r *Room) Join(m Member) (*roomMember, error) {
 		Member:  m,
 		Ignored: set.New(),
 	}
-	err := r.Members.Add(set.Itemize(m.ID(), member))
+	err := r.Members.AddNew(set.Itemize(m.ID(), member))
 	if err != nil {
 		return nil, err
 	}
diff --git a/host.go b/host.go
index 117af55..dfc7d55 100644
--- a/host.go
+++ b/host.go
@@ -187,6 +187,9 @@ func (h *Host) addClient(user *message.User, conn sshd.Connection) *Client {
 		timestamp: time.Now(),
 	}
 	h.mu.Lock()
+	if _, ok := h.clients[user]; ok {
+		logger.Warningf("user collision: %q", user)
+	}
 	h.clients[user] = append(h.clients[user], client)
 	h.mu.Unlock()
 	return &client
diff --git a/host_test.go b/host_test.go
index 784c521..2e93277 100644
--- a/host_test.go
+++ b/host_test.go
@@ -5,6 +5,7 @@ import (
 	"crypto/rand"
 	"crypto/rsa"
 	"errors"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"strings"
@@ -63,6 +64,8 @@ func TestHostNameCollision(t *testing.T) {
 
 	done := make(chan struct{}, 1)
 
+	canary := "canarystring"
+
 	// First client
 	go func() {
 		err := sshd.ConnectShell(s.Addr().String(), "foo", func(r io.Reader, w io.WriteCloser) error {
@@ -92,8 +95,10 @@ func TestHostNameCollision(t *testing.T) {
 				t.Errorf("Got %q; expected %q", actual, expected)
 			}
 
+			fmt.Fprintf(w, canary+message.Newline)
+
+			<-done
 			// Wrap it up.
-			close(done)
 			return nil
 		})
 		if err != nil {
@@ -108,16 +113,23 @@ func TestHostNameCollision(t *testing.T) {
 	err = sshd.ConnectShell(s.Addr().String(), "foo", func(r io.Reader, w io.WriteCloser) error {
 		scanner := bufio.NewScanner(r)
 
-		// Consume the initial buffer
-		scanner.Scan()
-		scanner.Scan()
-		scanner.Scan()
+		// Scan until we see our canarystring
+		for scanner.Scan() {
+			s := scanner.Text()
+			if strings.HasSuffix(s, canary) {
+				break
+			}
+		}
 
+		// Send an empty prompt to allow for a full line scan with EOL.
+		fmt.Fprintf(w, message.Newline)
+
+		scanner.Scan()
 		actual := scanner.Text()
 		if !strings.HasPrefix(actual, "[Guest1] ") {
-			// FIXME: Flaky?
 			t.Errorf("Second client did not get Guest1 name: %q", actual)
 		}
+		close(done)
 		return nil
 	})
 	if err != nil {