goreplay/output_ws_test.go

package goreplay

import (
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"sync"
	"testing"
)

func TestWebSocketOutput(t *testing.T) {
	wg := new(sync.WaitGroup)

	wsAddr := startWebsocket(func(data []byte) {
		wg.Done()
	})
	input := NewTestInput()
	output := NewWebSocketOutput(wsAddr, &WebSocketOutputConfig{Workers: 1})

	plugins := &InOutPlugins{
		Inputs:  []PluginReader{input},
		Outputs: []PluginWriter{output},
	}

	emitter := NewEmitter()
	go emitter.Start(plugins, Settings.Middleware)

	for i := 0; i < 10; i++ {
		wg.Add(1)
		input.EmitGET()
	}

	wg.Wait()
	emitter.Close()
}

func startWebsocket(cb func([]byte)) string {
	upgrader := websocket.Upgrader{}

	http.HandleFunc("/test", func(w http.ResponseWriter, r *http.Request) {
		c, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Print("upgrade:", err)
			return
		}

		go func(conn *websocket.Conn) {
			defer conn.Close()
			for {
				_, msg, _ := conn.ReadMessage()
				cb(msg)
			}
		}(c)
	})

	go func() {
		err := http.ListenAndServe("localhost:8081", nil)
		if err != nil {
			log.Fatal("Can't start:", err)
		}
	}()

	return "ws://localhost:8081/test"
}