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