1package terminal 2 3import ( 4 "bytes" 5 "net/http" 6 "net/http/httptest" 7 "testing" 8 9 "github.com/gorilla/websocket" 10 "github.com/stretchr/testify/assert" 11 "github.com/stretchr/testify/require" 12) 13 14const StreamMessage = "this is a test" 15 16func TestProxyStream(t *testing.T) { 17 downstream := bufferCloser{} 18 downstream.Write([]byte(StreamMessage)) 19 20 srv := streamServer{ 21 downstream: downstream, 22 } 23 24 s := httptest.NewServer(&srv) 25 defer s.Close() 26 27 c, _, err := websocket.DefaultDialer.Dial("ws://"+s.Listener.Addr().String()+"/ws", nil) 28 require.NoError(t, err) 29 30 // Check if writing to websocket works 31 c.WriteMessage(websocket.BinaryMessage, []byte(StreamMessage)) 32 b := make([]byte, len(StreamMessage)) 33 _, err = downstream.Read(b) 34 require.NoError(t, err) 35 assert.Equal(t, []byte(StreamMessage), b) 36 37 // Check if reading from websocket works 38 typ, b, err := c.ReadMessage() 39 require.NoError(t, err) 40 assert.Equal(t, typ, websocket.BinaryMessage) 41 assert.Equal(t, []byte(StreamMessage), b) 42} 43 44type streamServer struct { 45 downstream bufferCloser 46} 47 48func (d *streamServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { 49 proxy := NewStreamProxy(1) 50 ProxyStream(w, r, &d.downstream, proxy) 51} 52 53type bufferCloser struct { 54 bytes.Buffer 55} 56 57func (b *bufferCloser) Close() error { 58 return nil 59} 60