github.com/m-lab/locate@v0.17.6/connection/testdata/fakeserver.go (about)

     1  package testdata
     2  
     3  import (
     4  	"net/http"
     5  	"net/http/httptest"
     6  	"strings"
     7  	"sync"
     8  
     9  	"github.com/gorilla/websocket"
    10  )
    11  
    12  func FakeServer(handler func(http.ResponseWriter, *http.Request)) *httptest.Server {
    13  	mux := http.NewServeMux()
    14  	mux.Handle("/v2/heartbeat", http.HandlerFunc(handler))
    15  	s := httptest.NewServer(mux)
    16  	s.URL = strings.Replace(s.URL, "http", "ws", 1) + "/v2/heartbeat"
    17  	return s
    18  }
    19  
    20  type FakeHandler struct {
    21  	mu   sync.Mutex
    22  	conn *websocket.Conn
    23  }
    24  
    25  func (fh *FakeHandler) Upgrade(w http.ResponseWriter, r *http.Request) {
    26  	fh.mu.Lock()
    27  	defer fh.mu.Unlock()
    28  	upgrader := websocket.Upgrader{}
    29  	fh.conn, _ = upgrader.Upgrade(w, r, nil)
    30  }
    31  
    32  func (fh *FakeHandler) BadUpgrade(w http.ResponseWriter, r *http.Request) {
    33  	w.WriteHeader(400)
    34  }
    35  
    36  func (fh *FakeHandler) Read() ([]byte, error) {
    37  	fh.mu.Lock()
    38  	defer fh.mu.Unlock()
    39  	_, msg, err := fh.conn.ReadMessage()
    40  	return msg, err
    41  }
    42  
    43  func (fh *FakeHandler) Close() {
    44  	fh.mu.Lock()
    45  	defer fh.mu.Unlock()
    46  	fh.conn.Close()
    47  }