simplify tests

Author: hariso

pkg/foundation/grpcutil/websocket_test.go

@@ -15,6 +15,7 @@
 package grpcutil
 
 import (
+	"context"
 	"io"
 	"net/http"
 	"net/http/httptest"
@@ -27,56 +28,61 @@ import (
 )
 
 type testHandler struct {
-	responses chan string
+	is       *is.I
+	response string
 }
 
-func newTestHandler() *testHandler {
-	return &testHandler{responses: make(chan string, 1)}
-}
-
-func (h *testHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	select {
-	case resp := <-h.responses:
-		w.Write([]byte(resp))
-	}
+func (h *testHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
+	_, err := w.Write([]byte(h.response))
+	h.is.NoErr(err)
 }
 
 func TestWebSocket_NoUpgradeToWebSocket(t *testing.T) {
 	is := is.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
-	h := newTestHandler()
+	h := &testHandler{
+		is:       is,
+		response: "hi there",
+	}
 	s := httptest.NewServer(newWebSocketProxy(h, log.Nop()))
 	defer s.Close()
 
-	h.responses <- "hi there"
-	resp, err := http.Get(s.URL)
+	req, err := http.NewRequestWithContext(ctx, "GET", s.URL, nil)
+	is.NoErr(err)
+
+	resp, err := http.DefaultClient.Do(req)
 	is.NoErr(err)
 	is.True(resp.Body != nil) // expected response to have a body
 	defer resp.Body.Close()
 
 	bytes, err := io.ReadAll(resp.Body)
 	is.NoErr(err)
-	is.Equal("hi there", string(bytes))
+	is.Equal(h.response, string(bytes))
 }
 
 func TestWebSocket_UpgradeToWebSocket(t *testing.T) {
 	is := is.New(t)
 
-	h := newTestHandler()
+	h := &testHandler{
+		is:       is,
+		response: "hi there",
+	}
 	s := httptest.NewServer(newWebSocketProxy(h, log.Nop()))
 	defer s.Close()
 
 	// Convert http to ws
 	wsURL := "ws" + strings.TrimPrefix(s.URL, "http")
 
 	// Connect to the server
-	ws, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
+	ws, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
 	is.NoErr(err)
 	defer ws.Close()
+	defer resp.Body.Close()
 
-	h.responses <- "hi there"
 	msgType, bytes, err := ws.ReadMessage()
 	is.NoErr(err)
-	is.Equal("hi there", string(bytes))
+	is.Equal(h.response, string(bytes))
 	is.Equal(websocket.TextMessage, msgType)
 }