// new websocket connection wrapper func NewWSConnection(wsConn *websocket.Conn) *WSConnection { return &WSConnection{ id: wsConn.RemoteAddr().String(), wsConn: wsConn, writeChan: make(chan WSResponse, WriteChanBufferSize), // buffered. we keep track when its full quitChan: make(chan struct{}), } }
// new websocket connection wrapper func NewWSConnection(baseConn *websocket.Conn, funcMap map[string]*RPCFunc, evsw *events.EventSwitch) *WSConnection { wsc := &WSConnection{ id: baseConn.RemoteAddr().String(), baseConn: baseConn, writeChan: make(chan RPCResponse, writeChanCapacity), // error when full. funcMap: funcMap, evsw: evsw, } wsc.QuitService = *NewQuitService(log, "WSConnection", wsc) return wsc }
// unsubscribe from an event func unsubscribe(t *testing.T, con *websocket.Conn, eventid string) { err := con.WriteJSON(rpctypes.RPCRequest{ JSONRPC: "2.0", ID: "", Method: "unsubscribe", Params: []interface{}{eventid}, }) if err != nil { t.Fatal(err) } }
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 } } } }
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 } } }
// 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. } }