Go websocket test acting strange - go

Question

So basically I'm writing a go test for my chat application and for some reason the if I write Test_saveMessage function in the top of this file my tests go through and they work fine, however if I write the Test_InitRouter in the top of this file - my server opens and the test doesn't finish. As if it would be listening for more requests. Does anyone know the reason of why this could be happening? Here is the that does not work code:
package messenger
import (
"fmt"
"github.com/gorilla/websocket"
"github.com/stretchr/testify/assert"
"net/http/httptest"
"strings"
"testing"
)
var testMessage = Message{
Username: "Name",
Message: "Test message"}
//Tests InitRouter both sending and receiving messages
func Test_InitRouter(t *testing.T) {
var receivedMessage Message
//Create test server with the InitRouter handler
s := httptest.NewServer(InitRouter())
defer s.Close()
// Convert URL from http to ws
u := "ws" + strings.TrimPrefix(s.URL, "http")
fmt.Println(u)
// Connect to the test server
ws, _, err := websocket.DefaultDialer.Dial(u, nil)
if err != nil {
t.Fatalf("%v", err)
}
defer ws.Close()
//Send message to the server read received message and see if it's the same
if err != ws.WriteJSON(testMessage) {
t.Fatalf("%v", err)
}
err = ws.ReadJSON(&receivedMessage)
if err != nil {
t.Fatalf("%v", err)
}
if receivedMessage != testMessage {
t.Fatalf("%v", err)
}
}
//Test for the saveMessage function
func Test_saveMessage(t *testing.T) {
saveMessage(testMessage)
assert.Equal(t, 1, len(messages), "Expected to have 1 message")
}
As soon as I move the Test_saveMessage function to the top it starts working properly.
Here is the code for the handler:
package messenger
import (
"fmt"
"github.com/go-chi/chi"
"github.com/gorilla/websocket"
log "github.com/sirupsen/logrus"
"net/http"
)
func InitRouter() http.Handler {
r := chi.NewRouter()
r.Get("/", GetWebsocket)
return r
}
var clients = make(map[*websocket.Conn]bool) // connected clients
var broadcast = make(chan Message) // broadcast channel
var messages = []Message{}
func GetWebsocket(w http.ResponseWriter, r *http.Request) {
// Upgrade initial GET request to a websocket
upgrader := websocket.Upgrader{}
ws, err := upgrader.Upgrade(w, r, nil)
if err != nil {
log.Error(err)
}
// Close the connection when the function returns
defer ws.Close()
// Register our new client and send him the chat history
clients[ws] = true
serveInitialMessages(ws)
//initialize message sending logic
sendMessages(ws)
}
// Sends messages from a particular websocket to the channel
func sendMessages(ws *websocket.Conn){
for {
var msg Message
// Read in a new message as JSON and map it to a Message object
err := ws.ReadJSON(&msg)
if err != nil {
log.Info(err)
delete(clients, ws)
break
}
// Send the newly received message to the broadcast channel
broadcast <- msg
saveMessage(msg)
}
}
func HandleMessages() {
for {
// Grab the next message from the broadcast channel
msg := <-broadcast
fmt.Println(msg)
// Send it out to every client that is currently connected
for client := range clients {
err := client.WriteJSON(msg)
if err != nil {
log.Printf("error: %v", err)
client.Close()
delete(clients, client)
}
}
}
}
func saveMessage(m Message) {
if len(messages) >= 50 {
messages = messages[1:]
}
messages = append(messages, m)
}
func serveInitialMessages(ws *websocket.Conn) {
for _, m := range messages {
err := ws.WriteJSON(m)
if err != nil {
log.Error(err)
}
}
}