func reader(ws *websocket.Conn) { defer ws.Close() ws.SetReadLimit(512) ws.SetReadDeadline(time.Now().Add(pongWait)) ws.SetPongHandler(func(string) error { ws.SetReadDeadline(time.Now().Add(pongWait)); return nil }) for { _, _, err := ws.ReadMessage() if err != nil { break } } }
func writer(ws *websocket.Conn, lastMod time.Time) { lastError := "" pingTicker := time.NewTicker(pingPeriod) fileTicker := time.NewTicker(filePeriod) defer func() { pingTicker.Stop() fileTicker.Stop() ws.Close() }() for { select { case <-fileTicker.C: var p []byte var err error p, lastMod, err = readFileIfModified(lastMod) if err != nil { if s := err.Error(); s != lastError { lastError = s p = []byte(lastError) } } else { lastError = "" } if p != nil { ws.SetWriteDeadline(time.Now().Add(writeWait)) if err := ws.WriteMessage(websocket.TextMessage, p); err != nil { return } } case <-pingTicker.C: ws.SetWriteDeadline(time.Now().Add(writeWait)) if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil { return } } } }
// wait for an event; do things that might trigger events, and check them when they are received // the check function takes an event id and the byte slice read off the ws func waitForEvent(t *testing.T, con *websocket.Conn, eventid string, dieOnTimeout bool, f func(), check func(string, []byte) error) { // go routine to wait for webscoket msg goodCh := make(chan []byte) errCh := make(chan error) quitCh := make(chan struct{}) defer close(quitCh) // Read message go func() { for { _, p, err := con.ReadMessage() if err != nil { errCh <- err break } else { // if the event id isnt what we're waiting on // ignore it var response ctypes.Response var err error wire.ReadJSON(&response, p, &err) if err != nil { errCh <- err break } event, ok := response.Result.(*ctypes.ResultEvent) if ok && event.Event == eventid { goodCh <- p break } } } }() // do stuff (transactions) f() // wait for an event or timeout timeout := time.NewTimer(10 * time.Second) select { case <-timeout.C: if dieOnTimeout { con.Close() t.Fatalf("%s event was not received in time", eventid) } // else that's great, we didn't hear the event // and we shouldn't have case p := <-goodCh: if dieOnTimeout { // message was received and expected // run the check err := check(eventid, p) if err != nil { t.Fatal(err) panic(err) // Show the stack trace. } } else { con.Close() t.Fatalf("%s event was not expected", eventid) } case err := <-errCh: t.Fatal(err) panic(err) // Show the stack trace. } }