github.com/m-lab/locate@v0.17.6/handler/heartbeat_test.go (about)

     1  package handler
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/m-lab/locate/clientgeo"
    12  	"github.com/m-lab/locate/connection/testdata"
    13  	"github.com/m-lab/locate/heartbeat"
    14  	"github.com/m-lab/locate/heartbeat/heartbeattest"
    15  	prom "github.com/prometheus/client_golang/api/prometheus/v1"
    16  )
    17  
    18  func TestClient_Heartbeat_Error(t *testing.T) {
    19  	rw := httptest.NewRecorder()
    20  	// The header from this request will not contain the
    21  	// necessary "upgrade" tokens.
    22  	req := httptest.NewRequest(http.MethodGet, "/v2/heartbeat", nil)
    23  	c := fakeClient(nil)
    24  	c.Heartbeat(rw, req)
    25  
    26  	if rw.Code != http.StatusBadRequest {
    27  		t.Errorf("Heartbeat() wrong status code; got %d, want %d", rw.Code, http.StatusBadRequest)
    28  	}
    29  }
    30  
    31  func TestClient_handleHeartbeats(t *testing.T) {
    32  	wantErr := errors.New("connection error")
    33  	tests := []struct {
    34  		name    string
    35  		ws      conn
    36  		tracker heartbeat.StatusTracker
    37  	}{
    38  		{
    39  			name: "read-err",
    40  			ws: &fakeConn{
    41  				err: wantErr,
    42  			},
    43  		},
    44  		{
    45  			name: "registration-err",
    46  			ws: &fakeConn{
    47  				msg: testdata.FakeRegistration,
    48  			},
    49  			tracker: &heartbeattest.FakeStatusTracker{Err: wantErr},
    50  		},
    51  		{
    52  			name: "health-err",
    53  			ws: &fakeConn{
    54  				msg: testdata.FakeHealth,
    55  			},
    56  			tracker: &heartbeattest.FakeStatusTracker{Err: wantErr},
    57  		},
    58  	}
    59  	for _, tt := range tests {
    60  		t.Run(tt.name, func(t *testing.T) {
    61  			c := fakeClient(tt.tracker)
    62  			err := c.handleHeartbeats(tt.ws)
    63  			if !errors.Is(err, wantErr) {
    64  				t.Errorf("Client.handleHeartbeats() error = %v, wantErr %v", err, wantErr)
    65  			}
    66  		})
    67  	}
    68  }
    69  
    70  func fakeClient(t heartbeat.StatusTracker) *Client {
    71  	locatorv2 := fakeLocatorV2{StatusTracker: t}
    72  	return NewClient("mlab-sandbox", &fakeSigner{}, &locatorv2,
    73  		clientgeo.NewAppEngineLocator(), prom.NewAPI(nil), nil, nil, nil)
    74  }
    75  
    76  type fakeConn struct {
    77  	msg any
    78  	err error
    79  }
    80  
    81  // ReadMessage returns 0, the JSON encoding of a fake message, and an error.
    82  func (c *fakeConn) ReadMessage() (int, []byte, error) {
    83  	jsonMsg, _ := json.Marshal(c.msg)
    84  	return 0, jsonMsg, c.err
    85  }
    86  
    87  // SetReadDeadline returns nil.
    88  func (c *fakeConn) SetReadDeadline(time.Time) error {
    89  	return nil
    90  }
    91  
    92  // Close returns nil.
    93  func (c *fakeConn) Close() error {
    94  	return nil
    95  }