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  }