From 01b989cab2ff515c33a8153cf9e3a5a4c2ec0f65 Mon Sep 17 00:00:00 2001
From: Andrey Petrov <andrey.petrov@shazow.net>
Date: Thu, 8 Sep 2016 14:09:43 -0400
Subject: [PATCH] chat/message/user: screen -> WriteCloser embedded

---
 chat/message/user.go      | 18 ++++++++----------
 chat/message/user_test.go |  9 ++++++---
 chat/set_test.go          |  4 ++--
 3 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/chat/message/user.go b/chat/message/user.go
index d10d02b..fd656b2 100644
--- a/chat/message/user.go
+++ b/chat/message/user.go
@@ -18,13 +18,13 @@ var ErrUserClosed = errors.New("user closed")
 
 // User definition, implemented set Item interface and io.Writer
 type User struct {
-	colorIdx int
-	joined   time.Time
-	msg      chan Message
-	done     chan struct{}
+	io.WriteCloser
 
-	screen    io.WriteCloser
+	colorIdx  int
+	joined    time.Time
 	closeOnce sync.Once
+	msg       chan Message
+	done      chan struct{}
 
 	mu      sync.Mutex
 	name    string
@@ -47,7 +47,7 @@ func NewUser(name string) *User {
 
 func NewUserScreen(name string, screen io.WriteCloser) *User {
 	u := NewUser(name)
-	u.screen = screen
+	u.WriteCloser = screen
 
 	return u
 }
@@ -112,9 +112,7 @@ func (u *User) setColorIdx(idx int) {
 // Disconnect user, stop accepting messages
 func (u *User) Close() {
 	u.closeOnce.Do(func() {
-		if u.screen != nil {
-			u.screen.Close()
-		}
+		u.WriteCloser.Close()
 		// close(u.msg) TODO: Close?
 		close(u.done)
 	})
@@ -182,7 +180,7 @@ func (u *User) render(m Message) string {
 // HandleMsg will render the message to the screen, blocking.
 func (u *User) HandleMsg(m Message) error {
 	r := u.render(m)
-	_, err := u.screen.Write([]byte(r))
+	_, err := u.Write([]byte(r))
 	if err != nil {
 		logger.Printf("Write failed to %s, closing: %s", u.Name(), err)
 		u.Close()
diff --git a/chat/message/user_test.go b/chat/message/user_test.go
index 871730d..437015e 100644
--- a/chat/message/user_test.go
+++ b/chat/message/user_test.go
@@ -13,12 +13,15 @@ func TestMakeUser(t *testing.T) {
 	m := NewAnnounceMsg("hello")
 
 	defer u.Close()
-	u.Send(m)
-	u.HandleMsg(u.ConsumeOne())
+	err := u.Send(m)
+	if err != nil {
+		t.Fatalf("failed to send: %s", err)
+	}
+	u.HandleMsg(<-u.msg)
 
 	s.Read(&actual)
 	expected = []byte(m.String() + Newline)
 	if !reflect.DeepEqual(actual, expected) {
-		t.Errorf("Got: `%s`; Expected: `%s`", actual, expected)
+		t.Errorf("Got: %q; Expected: %q", actual, expected)
 	}
 }
diff --git a/chat/set_test.go b/chat/set_test.go
index 61722a8..754eb60 100644
--- a/chat/set_test.go
+++ b/chat/set_test.go
@@ -31,9 +31,9 @@ func TestSet(t *testing.T) {
 		t.Error(err)
 	}
 
-	err = s.Add(set.Itemize(u2.ID(), u2))
+	err = s.AddNew(set.Itemize(u2.ID(), u2))
 	if err != set.ErrCollision {
-		t.Error(err)
+		t.Errorf("expected ErrCollision, got: %s", err)
 	}
 
 	size := s.Len()