github.com/grahambrereton-form3/tilt@v0.10.18/internal/hud/server/websocket_test.go (about)

     1  package server
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/stretchr/testify/assert"
    10  
    11  	"github.com/windmilleng/tilt/internal/testutils"
    12  
    13  	"github.com/windmilleng/tilt/internal/store"
    14  )
    15  
    16  func TestWebsocketCloseOnReadErr(t *testing.T) {
    17  	ctx, _, _ := testutils.CtxAndAnalyticsForTest()
    18  	st, _ := store.NewStoreForTesting()
    19  	st.SetUpSubscribersForTesting(ctx)
    20  
    21  	conn := newFakeConn()
    22  	ws := NewWebsocketSubscriber(conn)
    23  	st.AddSubscriber(ctx, ws)
    24  
    25  	done := make(chan bool)
    26  	go func() {
    27  		ws.Stream(ctx, st)
    28  		close(done)
    29  	}()
    30  
    31  	st.NotifySubscribers(ctx)
    32  	conn.AssertNextWriteMsg(t).Ack()
    33  
    34  	st.NotifySubscribers(ctx)
    35  	conn.AssertNextWriteMsg(t).Ack()
    36  
    37  	conn.readCh <- fmt.Errorf("read error")
    38  
    39  	conn.AssertClose(t, done)
    40  }
    41  
    42  func TestWebsocketReadErrDuringMsg(t *testing.T) {
    43  	ctx, _, _ := testutils.CtxAndAnalyticsForTest()
    44  	st, _ := store.NewStoreForTesting()
    45  	st.SetUpSubscribersForTesting(ctx)
    46  
    47  	conn := newFakeConn()
    48  	ws := NewWebsocketSubscriber(conn)
    49  	st.AddSubscriber(ctx, ws)
    50  
    51  	done := make(chan bool)
    52  	go func() {
    53  		ws.Stream(ctx, st)
    54  		close(done)
    55  	}()
    56  
    57  	st.NotifySubscribers(ctx)
    58  
    59  	m := conn.AssertNextWriteMsg(t)
    60  
    61  	// Send a read error, and make sure the connection
    62  	// doesn't close immediately.
    63  	conn.readCh <- fmt.Errorf("read error")
    64  	time.Sleep(10 * time.Millisecond)
    65  	assert.False(t, conn.closed)
    66  
    67  	// Finish the write
    68  	m.Ack()
    69  
    70  	conn.AssertClose(t, done)
    71  }
    72  
    73  type fakeConn struct {
    74  	// Write an error to this channel to stop the Read consumer
    75  	readCh chan error
    76  
    77  	// Consume messages written to this channel. The caller should Ack() to acknowledge receipt.
    78  	writeCh chan msg
    79  
    80  	closed bool
    81  }
    82  
    83  func newFakeConn() *fakeConn {
    84  	return &fakeConn{
    85  		readCh:  make(chan error),
    86  		writeCh: make(chan msg),
    87  	}
    88  }
    89  
    90  func (c *fakeConn) NextReader() (int, io.Reader, error) {
    91  	return 1, nil, <-c.readCh
    92  }
    93  
    94  func (c *fakeConn) Close() error {
    95  	c.closed = true
    96  	return nil
    97  }
    98  
    99  func (c *fakeConn) WriteJSON(v interface{}) error {
   100  	msg := msg{callback: make(chan error)}
   101  	c.writeCh <- msg
   102  	return <-msg.callback
   103  }
   104  
   105  func (c *fakeConn) AssertNextWriteMsg(t *testing.T) msg {
   106  	select {
   107  	case <-time.After(100 * time.Millisecond):
   108  		t.Fatal("timed out waiting for WriteJSON")
   109  	case msg := <-c.writeCh:
   110  		return msg
   111  	}
   112  	return msg{}
   113  }
   114  
   115  func (c *fakeConn) AssertClose(t *testing.T, done chan bool) {
   116  	select {
   117  	case <-time.After(100 * time.Millisecond):
   118  		t.Fatal("timed out waiting for close")
   119  	case <-done:
   120  		assert.True(t, c.closed)
   121  	}
   122  }
   123  
   124  func (c *fakeConn) NextWriter(messagetype int) (io.WriteCloser, error) {
   125  	return c.writer(), nil
   126  }
   127  
   128  func (c *fakeConn) writer() io.WriteCloser {
   129  	return &fakeConnWriter{c: c}
   130  }
   131  
   132  type fakeConnWriter struct {
   133  	c *fakeConn
   134  }
   135  
   136  func (f *fakeConnWriter) Write(p []byte) (int, error) {
   137  	return len(p), nil
   138  }
   139  
   140  func (f *fakeConnWriter) Close() error {
   141  	cb := make(chan error)
   142  	f.c.writeCh <- msg{callback: cb}
   143  	return <-cb
   144  }
   145  
   146  type msg struct {
   147  	callback chan error
   148  }
   149  
   150  func (m msg) Ack() {
   151  	m.callback <- nil
   152  	close(m.callback)
   153  }