github.com/m-lab/locate@v0.17.6/cmd/heartbeat/main_test.go (about)

     1  package main
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"flag"
     7  	"net/url"
     8  	"os"
     9  	"reflect"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/m-lab/go/rtx"
    14  	v2 "github.com/m-lab/locate/api/v2"
    15  	"github.com/m-lab/locate/connection"
    16  	"github.com/m-lab/locate/connection/testdata"
    17  )
    18  
    19  func Test_main(t *testing.T) {
    20  	mainCtx, mainCancel = context.WithCancel(context.Background())
    21  	fh := testdata.FakeHandler{}
    22  	s := testdata.FakeServer(fh.Upgrade)
    23  	defer s.Close()
    24  	u, err := url.Parse(s.URL)
    25  	rtx.Must(err, "could not parse server URL")
    26  
    27  	lbPath = "/tmp/loadbalanced"
    28  	os.WriteFile(lbPath, []byte("false"), 0644)
    29  	defer os.Remove(lbPath)
    30  
    31  	flag.Set("heartbeat-url", s.URL)
    32  	flag.Set("hostname", "ndt-mlab1-lga0t.mlab-sandbox.measurement-lab.org")
    33  	flag.Set("experiment", "ndt")
    34  	flag.Set("pod", "ndt-abcde")
    35  	flag.Set("node", "mlab0-lga00.mlab-sandbox.measurement-lab.org")
    36  	flag.Set("namespace", "default")
    37  	flag.Set("registration-url", "file:./registration/testdata/registration.json")
    38  	flag.Set("services", "ndt/ndt7=ws://:"+u.Port()+"/ndt/v7/download")
    39  
    40  	heartbeatPeriod = 2 * time.Second
    41  	timer := time.NewTimer(2 * heartbeatPeriod)
    42  	go func() {
    43  		<-timer.C
    44  		if heartbeatURL != s.URL {
    45  			t.Errorf("main() incorrect locate URL; got: %s, want: %s",
    46  				heartbeatURL, s.URL)
    47  		}
    48  
    49  		msg, err := fh.Read()
    50  		rtx.Must(err, "could not read registration message")
    51  		var hbm v2.HeartbeatMessage
    52  		err = json.Unmarshal(msg, &hbm)
    53  		rtx.Must(err, "could not unmarshal registration message")
    54  		if hbm.Registration == nil {
    55  			t.Errorf("main() did not send registration message")
    56  		}
    57  
    58  		msg, err = fh.Read()
    59  		rtx.Must(err, "could not read health message")
    60  		err = json.Unmarshal(msg, &hbm)
    61  		rtx.Must(err, "could not unmarshal health message")
    62  		if hbm.Health.Score != 1 {
    63  			t.Errorf("write() did not send healthy (Score: 1) message")
    64  		}
    65  
    66  		mainCancel()
    67  	}()
    68  
    69  	main()
    70  }
    71  
    72  func Test_sendMessage(t *testing.T) {
    73  	tests := []struct {
    74  		name        string
    75  		msg         v2.HeartbeatMessage
    76  		msgType     string
    77  		wantDialMsg interface{}
    78  	}{
    79  		{
    80  			name:        "registration-change-dial-msg",
    81  			msg:         v2.HeartbeatMessage{Registration: &v2.Registration{City: "changed"}},
    82  			msgType:     "registration",
    83  			wantDialMsg: v2.HeartbeatMessage{Registration: &v2.Registration{City: "changed"}},
    84  		},
    85  		{
    86  			name:        "health-no-change",
    87  			msg:         v2.HeartbeatMessage{Health: &v2.Health{}},
    88  			msgType:     "health",
    89  			wantDialMsg: nil,
    90  		},
    91  	}
    92  	for _, tt := range tests {
    93  		t.Run(tt.name, func(t *testing.T) {
    94  			ws := connection.NewConn()
    95  			defer ws.Close()
    96  
    97  			sendMessage(ws, tt.msg, tt.msgType)
    98  			if !reflect.DeepEqual(ws.DialMessage, tt.wantDialMsg) {
    99  				t.Errorf("sendMessage() error updating websocket dial message; got: %v, want: %v",
   100  					ws.DialMessage, tt.wantDialMsg)
   101  			}
   102  		})
   103  	}
   104  }