github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/websocket_test.go (about) 1 package writer 2 3 import ( 4 "net/http" 5 "net/http/httptest" 6 "net/url" 7 "sync" 8 "testing" 9 "time" 10 11 "github.com/Jeffail/benthos/v3/lib/log" 12 "github.com/Jeffail/benthos/v3/lib/message" 13 "github.com/Jeffail/benthos/v3/lib/metrics" 14 "github.com/gorilla/websocket" 15 ) 16 17 func TestWebsocketBasic(t *testing.T) { 18 expMsgs := []string{ 19 "foo", 20 "bar", 21 "baz", 22 } 23 24 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 25 upgrader := websocket.Upgrader{} 26 27 var ws *websocket.Conn 28 var err error 29 if ws, err = upgrader.Upgrade(w, r, nil); err != nil { 30 return 31 } 32 33 defer ws.Close() 34 35 var actBytes []byte 36 for _, exp := range expMsgs { 37 if _, actBytes, err = ws.ReadMessage(); err != nil { 38 t.Error(err) 39 } else if act := string(actBytes); act != exp { 40 t.Errorf("Wrong msg contents: %v != %v", act, exp) 41 } 42 } 43 })) 44 45 conf := NewWebsocketConfig() 46 if wsURL, err := url.Parse(server.URL); err != nil { 47 t.Fatal(err) 48 } else { 49 wsURL.Scheme = "ws" 50 conf.URL = wsURL.String() 51 } 52 53 m, err := NewWebsocket(conf, log.Noop(), metrics.Noop()) 54 if err != nil { 55 t.Fatal(err) 56 } 57 58 if err = m.Connect(); err != nil { 59 t.Fatal(err) 60 } 61 62 for _, msg := range expMsgs { 63 if err = m.Write(message.New([][]byte{[]byte(msg)})); err != nil { 64 t.Error(err) 65 } 66 } 67 68 m.CloseAsync() 69 if err = m.WaitForClose(time.Second); err != nil { 70 t.Error(err) 71 } 72 } 73 74 func TestWebsocketClose(t *testing.T) { 75 closeChan := make(chan struct{}) 76 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 77 upgrader := websocket.Upgrader{} 78 79 var ws *websocket.Conn 80 var err error 81 if ws, err = upgrader.Upgrade(w, r, nil); err != nil { 82 return 83 } 84 85 ws.Close() 86 })) 87 88 conf := NewWebsocketConfig() 89 if wsURL, err := url.Parse(server.URL); err != nil { 90 t.Fatal(err) 91 } else { 92 wsURL.Scheme = "ws" 93 conf.URL = wsURL.String() 94 } 95 96 m, err := NewWebsocket(conf, log.Noop(), metrics.Noop()) 97 if err != nil { 98 t.Fatal(err) 99 } 100 101 if err = m.Connect(); err != nil { 102 t.Fatal(err) 103 } 104 105 wg := sync.WaitGroup{} 106 wg.Add(1) 107 go func() { 108 m.CloseAsync() 109 if cErr := m.WaitForClose(time.Second); cErr != nil { 110 t.Error(cErr) 111 } 112 wg.Done() 113 }() 114 115 wg.Wait() 116 close(closeChan) 117 }