Fixing websocket auth then sending challenge (#5406)

1 files changed, 63 insertions, 0 deletions

diff --git a/api/websocket_test.go b/api/websocket_test.go
index 6b8937d81..ab2959b03 100644
--- a/api/websocket_test.go
+++ b/api/websocket_test.go
@@ -251,6 +251,69 @@ func TestWebSocketEvent(t *testing.T) {
 	}
 }
 
+func TestCreateDirectChannelWithSocket(t *testing.T) {
+	th := Setup().InitBasic()
+	Client := th.BasicClient
+	user2 := th.BasicUser2
+
+	users := make([]*model.User, 0)
+	users = append(users, user2)
+
+	for i := 0; i < 10; i++ {
+		users = append(users, th.CreateUser(Client))
+	}
+
+	WebSocketClient, err := th.CreateWebSocketClient()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer WebSocketClient.Close()
+	WebSocketClient.Listen()
+
+	time.Sleep(300 * time.Millisecond)
+	if resp := <-WebSocketClient.ResponseChannel; resp.Status != model.STATUS_OK {
+		t.Fatal("should have responded OK to authentication challenge")
+	}
+
+	wsr := <-WebSocketClient.EventChannel
+	if wsr.Event != model.WEBSOCKET_EVENT_HELLO {
+		t.Fatal("missing hello")
+	}
+
+	stop := make(chan bool)
+	count := 0
+
+	go func() {
+		for {
+			select {
+			case wsr := <-WebSocketClient.EventChannel:
+				if wsr.Event == model.WEBSOCKET_EVENT_DIRECT_ADDED {
+					count = count + 1
+				}
+
+			case <-stop:
+				return
+			}
+		}
+	}()
+
+	for _, user := range users {
+		time.Sleep(100 * time.Millisecond)
+		if _, err := Client.CreateDirectChannel(user.Id); err != nil {
+			t.Fatal("failed to create DM channel")
+		}
+	}
+
+	time.Sleep(5000 * time.Millisecond)
+
+	stop <- true
+
+	if count != len(users) {
+		t.Fatal("We didn't get the proper amount of direct_added messages")
+	}
+
+}
+
 func TestWebsocketOriginSecurity(t *testing.T) {
 	Setup().InitBasic()
 	url := "ws://localhost" + utils.Cfg.ServiceSettings.ListenAddress