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

     1  package heartbeat
     2  
     3  import (
     4  	"math"
     5  	"math/rand"
     6  	"net/url"
     7  	"reflect"
     8  	"sort"
     9  	"testing"
    10  
    11  	"github.com/m-lab/go/host"
    12  	v2 "github.com/m-lab/locate/api/v2"
    13  	"github.com/m-lab/locate/heartbeat/heartbeattest"
    14  )
    15  
    16  var (
    17  	// Test services.
    18  	validNDT7Services = map[string][]string{
    19  		"ndt/ndt7": {
    20  			"ws://ndt/v7/upload",
    21  			"ws://ndt/v7/download",
    22  			"wss://ndt/v7/upload",
    23  			"wss://ndt/v7/download",
    24  		},
    25  	}
    26  
    27  	// Test URLs.
    28  	NDT7Urls = []url.URL{
    29  		{
    30  			Scheme: "ws",
    31  			Host:   "ndt",
    32  			Path:   "/v7/upload",
    33  		},
    34  		{
    35  			Scheme: "ws",
    36  			Host:   "ndt",
    37  			Path:   "/v7/download",
    38  		},
    39  		{
    40  			Scheme: "wss",
    41  			Host:   "ndt",
    42  			Path:   "/v7/upload",
    43  		},
    44  		{
    45  			Scheme: "wss",
    46  			Host:   "ndt",
    47  			Path:   "/v7/download",
    48  		},
    49  	}
    50  
    51  	// Test instances.
    52  	virtualInstance1 = v2.HeartbeatMessage{
    53  		Registration: &v2.Registration{
    54  			City:          "New York",
    55  			CountryCode:   "US",
    56  			ContinentCode: "NA",
    57  			Experiment:    "ndt",
    58  			Hostname:      "ndt-mlab1-lga00.mlab-sandbox.measurement-lab.org",
    59  			Latitude:      40.7667,
    60  			Longitude:     -73.8667,
    61  			Machine:       "mlab1",
    62  			Metro:         "lga",
    63  			Project:       "mlab-sandbox",
    64  			Probability:   1.0,
    65  			Site:          "lga00",
    66  			Type:          "virtual",
    67  			Uplink:        "10g",
    68  			Services:      validNDT7Services,
    69  		},
    70  		Health: &v2.Health{Score: 1},
    71  	}
    72  	virtualInstance2 = v2.HeartbeatMessage{
    73  		Registration: &v2.Registration{
    74  			City:          "New York",
    75  			CountryCode:   "US",
    76  			ContinentCode: "NA",
    77  			Experiment:    "ndt",
    78  			Hostname:      "ndt-mlab2-lga00.mlab-sandbox.measurement-lab.org",
    79  			Latitude:      40.7667,
    80  			Longitude:     -73.8667,
    81  			Machine:       "mlab2",
    82  			Metro:         "lga",
    83  			Project:       "mlab-sandbox",
    84  			Probability:   1.0,
    85  			Site:          "lga00",
    86  			Type:          "virtual",
    87  			Uplink:        "10g",
    88  			Services:      validNDT7Services,
    89  		},
    90  		Health: &v2.Health{Score: 1},
    91  	}
    92  	physicalInstance = v2.HeartbeatMessage{
    93  		Registration: &v2.Registration{
    94  			City:          "Los Angeles",
    95  			CountryCode:   "US",
    96  			ContinentCode: "NA",
    97  			Experiment:    "ndt",
    98  			Hostname:      "ndt-mlab1-lax00.mlab-sandbox.measurement-lab.org",
    99  			Latitude:      33.9425,
   100  			Longitude:     -118.4072,
   101  			Machine:       "mlab1",
   102  			Metro:         "lax",
   103  			Project:       "mlab-sandbox",
   104  			Probability:   1.0,
   105  			Site:          "lax00",
   106  			Type:          "physical",
   107  			Uplink:        "10g",
   108  			Services:      validNDT7Services,
   109  		},
   110  		Health: &v2.Health{Score: 1},
   111  	}
   112  	autonodeInstance = v2.HeartbeatMessage{
   113  		Registration: &v2.Registration{
   114  			City:          "Council Bluffs",
   115  			CountryCode:   "US",
   116  			ContinentCode: "NA",
   117  			Experiment:    "ndt",
   118  			Hostname:      "ndt-oma396982-2248791f.foo.sandbox.measurement-lab.org",
   119  			Latitude:      41.3032,
   120  			Longitude:     -95.8941,
   121  			Machine:       "2248791f",
   122  			Metro:         "oma",
   123  			Project:       "mlab-sandbox",
   124  			Probability:   1.0,
   125  			Site:          "oma396982",
   126  			Type:          "virtual",
   127  			Uplink:        "10g",
   128  			Services:      validNDT7Services,
   129  		},
   130  		Health: &v2.Health{Score: 1},
   131  	}
   132  	weheInstance = v2.HeartbeatMessage{
   133  		Registration: &v2.Registration{
   134  			City:          "Portland",
   135  			CountryCode:   "US",
   136  			ContinentCode: "NA",
   137  			Experiment:    "wehe",
   138  			Hostname:      "wehe-mlab1-pdx00.mlab-sandbox.measurement-lab.org",
   139  			Latitude:      45.5886,
   140  			Longitude:     -122.5975,
   141  			Machine:       "mlab1",
   142  			Metro:         "pdx",
   143  			Project:       "mlab-sandbox",
   144  			Probability:   1.0,
   145  			Site:          "pdx00",
   146  			Type:          "physical",
   147  			Uplink:        "10g",
   148  			Services:      map[string][]string{"wehe/replay": {"wss://4443/v0/envelope/access"}},
   149  		},
   150  		Health: &v2.Health{Score: 1},
   151  	}
   152  
   153  	// Test sites.
   154  	virtualSite = site{
   155  		distance: 296.04366543852825,
   156  		registration: v2.Registration{
   157  			City:          "New York",
   158  			CountryCode:   "US",
   159  			ContinentCode: "NA",
   160  			Experiment:    "ndt",
   161  			Latitude:      40.7667,
   162  			Longitude:     -73.8667,
   163  			Metro:         "lga",
   164  			Project:       "mlab-sandbox",
   165  			Probability:   1.0,
   166  			Site:          "lga00",
   167  			Type:          "virtual",
   168  			Uplink:        "10g",
   169  			Services:      validNDT7Services,
   170  		},
   171  		machines: []machine{
   172  			{
   173  				name:   "mlab1-lga00.mlab-sandbox.measurement-lab.org",
   174  				host:   "ndt-mlab1-lga00.mlab-sandbox.measurement-lab.org",
   175  				health: v2.Health{Score: 1},
   176  			},
   177  			{
   178  				name:   "mlab2-lga00.mlab-sandbox.measurement-lab.org",
   179  				host:   "ndt-mlab2-lga00.mlab-sandbox.measurement-lab.org",
   180  				health: v2.Health{Score: 1},
   181  			},
   182  		},
   183  	}
   184  	physicalSite = site{
   185  		distance: 3838.617961615054,
   186  		registration: v2.Registration{
   187  			City:          "Los Angeles",
   188  			CountryCode:   "US",
   189  			ContinentCode: "NA",
   190  			Experiment:    "ndt",
   191  			Latitude:      33.9425,
   192  			Longitude:     -118.4072,
   193  			Metro:         "lax",
   194  			Project:       "mlab-sandbox",
   195  			Probability:   1.0,
   196  			Site:          "lax00",
   197  			Type:          "physical",
   198  			Uplink:        "10g",
   199  			Services:      validNDT7Services,
   200  		},
   201  		machines: []machine{
   202  			{
   203  				name:   "mlab1-lax00.mlab-sandbox.measurement-lab.org",
   204  				host:   "ndt-mlab1-lax00.mlab-sandbox.measurement-lab.org",
   205  				health: v2.Health{Score: 1},
   206  			},
   207  		},
   208  	}
   209  	autonodeSite = site{
   210  		distance: 1701.749354381346,
   211  		registration: v2.Registration{
   212  			City:          "Council Bluffs",
   213  			CountryCode:   "US",
   214  			ContinentCode: "NA",
   215  			Experiment:    "ndt",
   216  			Latitude:      41.3032,
   217  			Longitude:     -95.8941,
   218  			Metro:         "oma",
   219  			Project:       "mlab-sandbox",
   220  			Probability:   1.0,
   221  			Site:          "oma396982",
   222  			Type:          "virtual",
   223  			Uplink:        "10g",
   224  			Services:      validNDT7Services,
   225  		},
   226  		machines: []machine{
   227  			{
   228  				name:   "ndt-oma396982-2248791f.foo.sandbox.measurement-lab.org",
   229  				host:   "ndt-oma396982-2248791f.foo.sandbox.measurement-lab.org",
   230  				health: v2.Health{Score: 1},
   231  			},
   232  		},
   233  	}
   234  	weheSite = site{
   235  		distance: 3710.7679340078703,
   236  		registration: v2.Registration{
   237  			City:          "Portland",
   238  			CountryCode:   "US",
   239  			ContinentCode: "NA",
   240  			Experiment:    "wehe",
   241  			Latitude:      45.5886,
   242  			Longitude:     -122.5975,
   243  			Metro:         "pdx",
   244  			Project:       "mlab-sandbox",
   245  			Probability:   1.0,
   246  			Site:          "pdx00",
   247  			Type:          "physical",
   248  			Uplink:        "10g",
   249  			Services:      map[string][]string{"wehe/replay": {"wss://4443/v0/envelope/access"}},
   250  		},
   251  		machines: []machine{
   252  			{
   253  				name:   "mlab1-pdx00.mlab-sandbox.measurement-lab.org",
   254  				host:   "wehe-mlab1-pdx00.mlab-sandbox.measurement-lab.org",
   255  				health: v2.Health{Score: 1},
   256  			},
   257  		},
   258  	}
   259  
   260  	// Test Targets.
   261  	virtualTarget = v2.Target{
   262  		Machine:  "mlab1-lga00.mlab-sandbox.measurement-lab.org",
   263  		Hostname: "ndt-mlab1-lga00.mlab-sandbox.measurement-lab.org",
   264  		Location: &v2.Location{
   265  			City:    "New York",
   266  			Country: "US",
   267  		},
   268  		URLs: map[string]string{},
   269  	}
   270  	physicalTarget = v2.Target{
   271  		Machine:  "mlab1-lax00.mlab-sandbox.measurement-lab.org",
   272  		Hostname: "ndt-mlab1-lax00.mlab-sandbox.measurement-lab.org",
   273  		Location: &v2.Location{
   274  			City:    "Los Angeles",
   275  			Country: "US",
   276  		},
   277  		URLs: map[string]string{},
   278  	}
   279  	weheTarget = v2.Target{
   280  		Machine:  "mlab1-pdx00.mlab-sandbox.measurement-lab.org",
   281  		Hostname: "wehe-mlab1-pdx00.mlab-sandbox.measurement-lab.org",
   282  		Location: &v2.Location{
   283  			City:    "Portland",
   284  			Country: "US",
   285  		},
   286  		URLs: map[string]string{},
   287  	}
   288  )
   289  
   290  func TestNearest(t *testing.T) {
   291  	instances := []v2.HeartbeatMessage{
   292  		virtualInstance1,
   293  		physicalInstance,
   294  		weheInstance,
   295  	}
   296  
   297  	tests := []struct {
   298  		name      string
   299  		service   string
   300  		lat       float64
   301  		lon       float64
   302  		instances []v2.HeartbeatMessage
   303  		opts      *NearestOptions
   304  		expected  *TargetInfo
   305  		wantErr   bool
   306  	}{
   307  		{
   308  			// Test client coordinates are in NY, virtual target in LGA, and physical target in LAX.
   309  			name:    "NDT7-any-type",
   310  			service: "ndt/ndt7",
   311  			lat:     43.1988,
   312  			lon:     -75.3242,
   313  			opts:    &NearestOptions{Type: "", Country: "US"},
   314  			expected: &TargetInfo{
   315  				Targets: []v2.Target{virtualTarget, physicalTarget},
   316  				URLs:    NDT7Urls,
   317  				Ranks:   map[string]int{virtualTarget.Machine: 0, physicalTarget.Machine: 1},
   318  			},
   319  			wantErr: false,
   320  		},
   321  		{
   322  			name:    "NDT7-physical",
   323  			service: "ndt/ndt7",
   324  			lat:     43.1988,
   325  			lon:     -75.3242,
   326  			opts:    &NearestOptions{Type: "physical", Country: "US"},
   327  			expected: &TargetInfo{
   328  				Targets: []v2.Target{physicalTarget},
   329  				URLs:    NDT7Urls,
   330  				Ranks:   map[string]int{physicalTarget.Machine: 0},
   331  			},
   332  			wantErr: false,
   333  		},
   334  		{
   335  			name:    "NDT7-virtual",
   336  			service: "ndt/ndt7",
   337  			lat:     43.1988,
   338  			lon:     -75.3242,
   339  			opts:    &NearestOptions{Type: "virtual", Country: "US"},
   340  			expected: &TargetInfo{
   341  				Targets: []v2.Target{virtualTarget},
   342  				URLs:    NDT7Urls,
   343  				Ranks:   map[string]int{virtualTarget.Machine: 0},
   344  			},
   345  			wantErr: false,
   346  		},
   347  		{
   348  			name:    "wehe",
   349  			service: "wehe/replay",
   350  			lat:     43.1988,
   351  			lon:     -75.3242,
   352  			opts:    &NearestOptions{Type: "", Country: "US"},
   353  			expected: &TargetInfo{
   354  				Targets: []v2.Target{weheTarget},
   355  				URLs: []url.URL{{
   356  					Scheme: "wss",
   357  					Host:   "4443",
   358  					Path:   "/v0/envelope/access",
   359  				}},
   360  				Ranks: map[string]int{weheTarget.Machine: 0},
   361  			},
   362  			wantErr: false,
   363  		},
   364  		{
   365  			// Test client coordinates are in NY, virtual target in LGA, and physical target in LAX.
   366  			name:    "NDT-sites-found",
   367  			service: "ndt/ndt7",
   368  			lat:     43.1988,
   369  			lon:     -75.3242,
   370  			opts:    &NearestOptions{Type: "", Country: "US", Sites: []string{"lga00", "lax00"}},
   371  			expected: &TargetInfo{
   372  				Targets: []v2.Target{virtualTarget, physicalTarget},
   373  				URLs:    NDT7Urls,
   374  				Ranks:   map[string]int{virtualTarget.Machine: 0, physicalTarget.Machine: 1},
   375  			},
   376  			wantErr: false,
   377  		},
   378  		{
   379  			name:     "NDT-sites-empty",
   380  			service:  "ndt/ndt7",
   381  			lat:      43.1988,
   382  			lon:      -75.3242,
   383  			opts:     &NearestOptions{Type: "", Country: "US", Sites: []string{"foo99", "bar99"}},
   384  			expected: nil,
   385  			wantErr:  true,
   386  		},
   387  		{
   388  			name:    "NDT7-any-type-country",
   389  			service: "ndt/ndt7",
   390  			lat:     43.1988,
   391  			lon:     -75.3242,
   392  			opts:    &NearestOptions{Type: "", Country: "IT"},
   393  			expected: &TargetInfo{
   394  				Targets: []v2.Target{virtualTarget, physicalTarget},
   395  				URLs:    NDT7Urls,
   396  				Ranks:   map[string]int{virtualTarget.Machine: 0, physicalTarget.Machine: 1},
   397  			},
   398  			wantErr: false,
   399  		},
   400  		{
   401  			name:     "NDT7-any-type-country-strict",
   402  			service:  "ndt/ndt7",
   403  			lat:      43.1988,
   404  			lon:      -75.3242,
   405  			opts:     &NearestOptions{Type: "", Country: "IT", Strict: true},
   406  			expected: nil,
   407  			wantErr:  true,
   408  		},
   409  	}
   410  
   411  	for _, tt := range tests {
   412  		t.Run(tt.name, func(t *testing.T) {
   413  			memorystore := heartbeattest.FakeMemorystoreClient
   414  			tracker := NewHeartbeatStatusTracker(&memorystore)
   415  			locator := NewServerLocator(tracker)
   416  			locator.StopImport()
   417  			rand.Seed(1658458451000000000)
   418  
   419  			for _, i := range instances {
   420  				locator.RegisterInstance(*i.Registration)
   421  				locator.UpdateHealth(i.Registration.Hostname, *i.Health)
   422  			}
   423  
   424  			got, err := locator.Nearest(tt.service, tt.lat, tt.lon, tt.opts)
   425  
   426  			if (err != nil) != tt.wantErr {
   427  				t.Fatalf("Nearest() error got: %t, want %t, err: %v", err != nil, tt.wantErr, err)
   428  			}
   429  
   430  			if !reflect.DeepEqual(got, tt.expected) {
   431  				t.Errorf("Nearest() targets got: %+v, want %+v", got, tt.expected)
   432  			}
   433  		})
   434  	}
   435  
   436  }
   437  
   438  func TestFilterSites(t *testing.T) {
   439  	instances := map[string]v2.HeartbeatMessage{
   440  		"virtual1": virtualInstance1,
   441  		"virtual2": virtualInstance2,
   442  		"physical": physicalInstance,
   443  		"autonode": autonodeInstance,
   444  		"wehe":     weheInstance,
   445  	}
   446  
   447  	tests := []struct {
   448  		name     string
   449  		service  string
   450  		typ      string
   451  		country  string
   452  		strict   bool
   453  		org      string
   454  		lat      float64
   455  		lon      float64
   456  		expected []site
   457  	}{
   458  		{
   459  			name:     "NDT7-any-type",
   460  			service:  "ndt/ndt7",
   461  			typ:      "",
   462  			country:  "US",
   463  			lat:      43.1988,
   464  			lon:      -75.3242,
   465  			expected: []site{virtualSite, autonodeSite, physicalSite},
   466  		},
   467  		{
   468  			name:     "NDT7-physical",
   469  			service:  "ndt/ndt7",
   470  			typ:      "physical",
   471  			country:  "US",
   472  			lat:      43.1988,
   473  			lon:      -75.3242,
   474  			expected: []site{physicalSite},
   475  		},
   476  		{
   477  			name:     "NDT7-virtual",
   478  			service:  "ndt/ndt7",
   479  			typ:      "virtual",
   480  			country:  "US",
   481  			lat:      43.1988,
   482  			lon:      -75.3242,
   483  			expected: []site{virtualSite, autonodeSite},
   484  		},
   485  		{
   486  			name:     "wehe",
   487  			service:  "wehe/replay",
   488  			typ:      "",
   489  			country:  "US",
   490  			lat:      43.1988,
   491  			lon:      -75.3242,
   492  			expected: []site{weheSite},
   493  		},
   494  		{
   495  			name:     "too-far",
   496  			service:  "ndt-ndt7",
   497  			typ:      "",
   498  			country:  "",
   499  			lat:      1000,
   500  			lon:      1000,
   501  			expected: []site{},
   502  		},
   503  		{
   504  			name:     "country-with-strict",
   505  			service:  "ndt/ndt7",
   506  			typ:      "",
   507  			country:  "US",
   508  			strict:   true,
   509  			lat:      43.1988,
   510  			lon:      -75.3242,
   511  			expected: []site{virtualSite, autonodeSite, physicalSite},
   512  		},
   513  		{
   514  			name:     "country-with-strict-no-results",
   515  			service:  "ndt/ndt7",
   516  			typ:      "",
   517  			country:  "IT",
   518  			strict:   true,
   519  			lat:      43.1988,
   520  			lon:      -75.3242,
   521  			expected: []site{},
   522  		},
   523  		{
   524  			name:     "org-skip-v2-names",
   525  			service:  "ndt/ndt7",
   526  			org:      "foo",
   527  			lat:      43.1988,
   528  			lon:      -75.3242,
   529  			expected: []site{autonodeSite},
   530  		},
   531  		{
   532  			name:     "org-skip-v3-names-different-org",
   533  			service:  "ndt/ndt7",
   534  			org:      "zoom",
   535  			lat:      43.1988,
   536  			lon:      -75.3242,
   537  			expected: []site{},
   538  		},
   539  		{
   540  			name:     "org-allow-v2-names-for-mlab-org",
   541  			service:  "ndt/ndt7",
   542  			org:      "mlab",
   543  			lat:      43.1988,
   544  			lon:      -75.3242,
   545  			expected: []site{virtualSite, physicalSite},
   546  		},
   547  	}
   548  
   549  	for _, tt := range tests {
   550  		t.Run(tt.name, func(t *testing.T) {
   551  			opts := &NearestOptions{Type: tt.typ, Country: tt.country, Strict: tt.strict, Org: tt.org}
   552  			got := filterSites(tt.service, tt.lat, tt.lon, instances, opts)
   553  
   554  			sortSites(got)
   555  			for _, v := range got {
   556  				sort.Slice(v.machines, func(i, j int) bool {
   557  					return v.machines[i].name < v.machines[j].name
   558  				})
   559  			}
   560  
   561  			if !reflect.DeepEqual(got, tt.expected) {
   562  				t.Errorf("filterSites()\n got: %+v\nwant: %+v", got, tt.expected)
   563  			}
   564  		})
   565  	}
   566  }
   567  
   568  func TestIsValidInstance(t *testing.T) {
   569  	validHost := "ndt-mlab1-lga00.mlab-sandbox.measurement-lab.org"
   570  	validLat := 40.7667
   571  	validLon := -73.8667
   572  	validType := "virtual"
   573  	validScore := float64(1)
   574  
   575  	tests := []struct {
   576  		name         string
   577  		typ          string
   578  		host         string
   579  		lat          float64
   580  		lon          float64
   581  		instanceType string
   582  		services     map[string][]string
   583  		score        float64
   584  		prom         *v2.Prometheus
   585  		expected     bool
   586  		expectedHost host.Name
   587  		expectedDist float64
   588  	}{
   589  		{
   590  			name:         "0-health",
   591  			typ:          "virtual",
   592  			host:         validHost,
   593  			lat:          validLat,
   594  			lon:          validLon,
   595  			services:     validNDT7Services,
   596  			instanceType: validType,
   597  			score:        0,
   598  			expected:     false,
   599  			expectedHost: host.Name{},
   600  			expectedDist: 0,
   601  		},
   602  		{
   603  			name:         "prometheus-unhealthy",
   604  			typ:          "",
   605  			host:         validHost,
   606  			lat:          validLat,
   607  			lon:          validLon,
   608  			services:     validNDT7Services,
   609  			instanceType: validType,
   610  			score:        validScore,
   611  			prom: &v2.Prometheus{
   612  				Health: false,
   613  			},
   614  			expected:     false,
   615  			expectedHost: host.Name{},
   616  			expectedDist: 0,
   617  		},
   618  		{
   619  			name:         "invalid-host",
   620  			typ:          "virtual",
   621  			host:         "invalid-host",
   622  			lat:          validLat,
   623  			lon:          validLon,
   624  			services:     validNDT7Services,
   625  			instanceType: validType,
   626  			score:        validScore,
   627  			expected:     false,
   628  			expectedHost: host.Name{},
   629  			expectedDist: 0,
   630  		},
   631  		{
   632  			name:         "mismatched-type",
   633  			typ:          "virtual",
   634  			host:         validHost,
   635  			lat:          validLat,
   636  			lon:          validLon,
   637  			services:     validNDT7Services,
   638  			instanceType: "physical",
   639  			score:        validScore,
   640  			expected:     false,
   641  			expectedHost: host.Name{},
   642  			expectedDist: 0,
   643  		},
   644  		{
   645  			name:         "invalid-service",
   646  			typ:          "virtual",
   647  			host:         validHost,
   648  			lat:          validLat,
   649  			lon:          validLon,
   650  			services:     map[string][]string{},
   651  			instanceType: validType,
   652  			score:        validScore,
   653  			expected:     false,
   654  			expectedHost: host.Name{},
   655  			expectedDist: 0,
   656  		},
   657  		{
   658  			name:         "success-same-type",
   659  			typ:          "virtual",
   660  			host:         validHost,
   661  			lat:          validLat,
   662  			lon:          validLon,
   663  			services:     validNDT7Services,
   664  			instanceType: validType,
   665  			score:        validScore,
   666  			expected:     true,
   667  			expectedHost: host.Name{
   668  				Service: "ndt",
   669  				Machine: "mlab1",
   670  				Site:    "lga00",
   671  				Project: "mlab-sandbox",
   672  				Domain:  "measurement-lab.org",
   673  				Suffix:  "",
   674  				Version: "v2",
   675  			},
   676  			expectedDist: 296.043665,
   677  		},
   678  		{
   679  			name:         "success-no-type",
   680  			typ:          "",
   681  			host:         validHost,
   682  			lat:          validLat,
   683  			lon:          validLon,
   684  			services:     validNDT7Services,
   685  			instanceType: validType,
   686  			score:        validScore,
   687  			expected:     true,
   688  			expectedHost: host.Name{
   689  				Service: "ndt",
   690  				Machine: "mlab1",
   691  				Site:    "lga00",
   692  				Project: "mlab-sandbox",
   693  				Domain:  "measurement-lab.org",
   694  				Suffix:  "",
   695  				Version: "v2",
   696  			},
   697  			expectedDist: 296.043665,
   698  		},
   699  	}
   700  	for _, tt := range tests {
   701  		t.Run(tt.name, func(t *testing.T) {
   702  			v := v2.HeartbeatMessage{
   703  				Registration: &v2.Registration{
   704  					City:          "New York",
   705  					CountryCode:   "US",
   706  					ContinentCode: "NA",
   707  					Experiment:    "ndt",
   708  					Hostname:      tt.host,
   709  					Latitude:      tt.lat,
   710  					Longitude:     tt.lon,
   711  					Machine:       "mlab1",
   712  					Metro:         "lga",
   713  					Project:       "mlab-sandbox",
   714  					Probability:   1.0,
   715  					Site:          "lga00",
   716  					Type:          tt.instanceType,
   717  					Uplink:        "10g",
   718  					Services:      tt.services,
   719  				},
   720  				Health: &v2.Health{
   721  					Score: tt.score,
   722  				},
   723  				Prometheus: tt.prom,
   724  			}
   725  			opts := &NearestOptions{Type: tt.typ}
   726  			got, gotHost, gotDist := isValidInstance("ndt/ndt7", 43.1988, -75.3242, v, opts)
   727  
   728  			if got != tt.expected {
   729  				t.Errorf("isValidInstance() got: %t, want: %t", got, tt.expected)
   730  			}
   731  
   732  			if gotHost != tt.expectedHost {
   733  				t.Errorf("isValidInstance() host got: %#v, want: %#v", gotHost, tt.expectedHost)
   734  			}
   735  
   736  			if math.Abs(gotDist-tt.expectedDist) > 0.01 {
   737  				t.Errorf("isValidInstance() distance got: %f, want: %f", gotDist, tt.expectedDist)
   738  			}
   739  		})
   740  	}
   741  }
   742  
   743  func TestSortSites(t *testing.T) {
   744  	tests := []struct {
   745  		name     string
   746  		sites    []site
   747  		expected []site
   748  	}{
   749  		{
   750  			name:     "empty",
   751  			sites:    []site{},
   752  			expected: []site{},
   753  		},
   754  		{
   755  			name:     "one",
   756  			sites:    []site{{distance: 10}},
   757  			expected: []site{{distance: 10}},
   758  		},
   759  		{
   760  			name: "many",
   761  			sites: []site{{distance: 3838.61}, {distance: 3710.7679340078703}, {distance: -895420.92},
   762  				{distance: 296.0436}, {distance: math.MaxFloat64}, {distance: 3838.61}},
   763  			expected: []site{{distance: -895420.92}, {distance: 296.0436}, {distance: 3710.7679340078703},
   764  				{distance: 3838.61}, {distance: 3838.61}, {distance: math.MaxFloat64}},
   765  		},
   766  	}
   767  
   768  	for _, tt := range tests {
   769  		t.Run(tt.name, func(t *testing.T) {
   770  			sortSites(tt.sites)
   771  
   772  			if !reflect.DeepEqual(tt.sites, tt.expected) {
   773  				t.Errorf("sortSites() got: %+v, want: %+v", tt.sites, tt.expected)
   774  			}
   775  		})
   776  	}
   777  }
   778  
   779  func TestRankSites(t *testing.T) {
   780  	tests := []struct {
   781  		name     string
   782  		sites    []site
   783  		expected []site
   784  	}{
   785  		{
   786  			name:     "empty",
   787  			sites:    []site{},
   788  			expected: []site{},
   789  		},
   790  		{
   791  			name:     "one",
   792  			sites:    []site{{distance: 10}},
   793  			expected: []site{{distance: 10, rank: 0, metroRank: 0}},
   794  		},
   795  		{
   796  			name: "many",
   797  			sites: []site{
   798  				{registration: v2.Registration{Metro: "a"}},
   799  				{registration: v2.Registration{Metro: "b"}},
   800  				{registration: v2.Registration{Metro: "b"}},
   801  				{registration: v2.Registration{Metro: "c"}},
   802  				{registration: v2.Registration{Metro: "b"}}},
   803  			expected: []site{
   804  				{rank: 0, metroRank: 0, registration: v2.Registration{Metro: "a"}},
   805  				{rank: 1, metroRank: 1, registration: v2.Registration{Metro: "b"}},
   806  				{rank: 2, metroRank: 1, registration: v2.Registration{Metro: "b"}},
   807  				{rank: 3, metroRank: 2, registration: v2.Registration{Metro: "c"}},
   808  				{rank: 4, metroRank: 1, registration: v2.Registration{Metro: "b"}},
   809  			},
   810  		},
   811  	}
   812  
   813  	for _, tt := range tests {
   814  		t.Run(tt.name, func(t *testing.T) {
   815  			rank(tt.sites)
   816  
   817  			if !reflect.DeepEqual(tt.sites, tt.expected) {
   818  				t.Errorf("rankSites() got: %+v, want: %+v", tt.sites, tt.expected)
   819  			}
   820  		})
   821  	}
   822  }
   823  
   824  func TestPickTargets(t *testing.T) {
   825  	// Sites numbered by distance, which makes it easier to understand expected values.
   826  	site1 := site{
   827  		distance: 10,
   828  		registration: v2.Registration{
   829  			City:        "New York",
   830  			CountryCode: "US",
   831  			Services:    validNDT7Services,
   832  			Metro:       "lga",
   833  		},
   834  		metroRank: 0,
   835  		machines: []machine{
   836  			{name: "mlab1-site1-metro0", host: "ndt-mlab1-site1-metro0"},
   837  			{name: "mlab2-site1-metro0", host: "ndt-mlab2-site1-metro0"},
   838  			{name: "mlab3-site1-metro0", host: "ndt-mlab3-site1-metro0"},
   839  			{name: "mlab4-site-metro10", host: "ndt-mlab4-site-metro10"},
   840  		},
   841  	}
   842  	site2 := site{
   843  		distance: 10,
   844  		registration: v2.Registration{
   845  			City:        "New York",
   846  			CountryCode: "US",
   847  			Services:    validNDT7Services,
   848  			Metro:       "lga",
   849  		},
   850  		metroRank: 0,
   851  		machines: []machine{
   852  			{name: "mlab1-site2-metro0", host: "ndt-mlab1-site2-metro0"},
   853  			{name: "mlab2-site2-metro0", host: "ndt-mlab2-site2-metro0"},
   854  			{name: "mlab3-site2-metro0", host: "ndt-mlab3-site2-metro0"},
   855  			{name: "mlab4-site2-metro0", host: "ndt-mlab4-site2-metro0"},
   856  		},
   857  	}
   858  	site3 := site{
   859  		distance: 100,
   860  		registration: v2.Registration{
   861  			City:        "Los Angeles",
   862  			CountryCode: "US",
   863  			Services:    validNDT7Services,
   864  			Metro:       "lax",
   865  		},
   866  		metroRank: 1,
   867  		machines: []machine{
   868  			{name: "mlab1-site3-metro1", host: "ndt-mlab1-site3-metro1"},
   869  		},
   870  	}
   871  	site4 := site{
   872  		distance: 110,
   873  		registration: v2.Registration{
   874  			City:        "Portland",
   875  			CountryCode: "US",
   876  			Services:    validNDT7Services,
   877  			Metro:       "pdx",
   878  		},
   879  		metroRank: 2,
   880  		machines: []machine{
   881  			{name: "mlab1-site4-metro2", host: "ndt-mlab1-site4-metro2"},
   882  		},
   883  	}
   884  
   885  	tests := []struct {
   886  		name     string
   887  		sites    []site
   888  		expected *TargetInfo
   889  	}{
   890  		{
   891  			name: "4-sites",
   892  			sites: []site{
   893  				site1, site2, site3, site4,
   894  			},
   895  			expected: &TargetInfo{
   896  				Targets: []v2.Target{
   897  					{
   898  						Machine:  "mlab2-site2-metro0",
   899  						Hostname: "ndt-mlab2-site2-metro0",
   900  						Location: &v2.Location{
   901  							City:    site2.registration.City,
   902  							Country: site2.registration.CountryCode,
   903  						},
   904  						URLs: make(map[string]string),
   905  					},
   906  					{
   907  						Machine:  "mlab3-site1-metro0",
   908  						Hostname: "ndt-mlab3-site1-metro0",
   909  						Location: &v2.Location{
   910  							City:    site1.registration.City,
   911  							Country: site1.registration.CountryCode,
   912  						},
   913  						URLs: make(map[string]string),
   914  					},
   915  					{
   916  						Machine:  "mlab1-site3-metro1",
   917  						Hostname: "ndt-mlab1-site3-metro1",
   918  						Location: &v2.Location{
   919  							City:    site3.registration.City,
   920  							Country: site3.registration.CountryCode,
   921  						},
   922  						URLs: make(map[string]string),
   923  					},
   924  					{
   925  						Machine:  "mlab1-site4-metro2",
   926  						Hostname: "ndt-mlab1-site4-metro2",
   927  						Location: &v2.Location{
   928  							City:    site4.registration.City,
   929  							Country: site4.registration.CountryCode,
   930  						},
   931  						URLs: make(map[string]string),
   932  					},
   933  				},
   934  				URLs: NDT7Urls,
   935  				Ranks: map[string]int{
   936  					"mlab1-site3-metro1": 1,
   937  					"mlab1-site4-metro2": 2,
   938  					"mlab2-site2-metro0": 0,
   939  					"mlab3-site1-metro0": 0,
   940  				},
   941  			},
   942  		},
   943  		{
   944  			name: "1-site",
   945  			sites: []site{
   946  				site1,
   947  			},
   948  			expected: &TargetInfo{
   949  				Targets: []v2.Target{
   950  					{
   951  						Machine:  "mlab2-site1-metro0",
   952  						Hostname: "ndt-mlab2-site1-metro0",
   953  						Location: &v2.Location{
   954  							City:    site1.registration.City,
   955  							Country: site1.registration.CountryCode,
   956  						},
   957  						URLs: make(map[string]string),
   958  					},
   959  				},
   960  				URLs:  NDT7Urls,
   961  				Ranks: map[string]int{"mlab2-site1-metro0": 0},
   962  			},
   963  		},
   964  	}
   965  	for _, tt := range tests {
   966  		t.Run(tt.name, func(t *testing.T) {
   967  			// Use a fixed seed so the pattern is only pseudorandom and can
   968  			// be verififed against expectations.
   969  			rand.Seed(1658340109320624212)
   970  			got := pickTargets("ndt/ndt7", tt.sites)
   971  
   972  			if !reflect.DeepEqual(got, tt.expected) {
   973  				t.Errorf("pickTargets() got: %+v, want: %+v", got, tt.expected)
   974  			}
   975  		})
   976  	}
   977  }
   978  
   979  func TestAlwaysPick(t *testing.T) {
   980  	tests := []struct {
   981  		name string
   982  		opts *NearestOptions
   983  		want bool
   984  	}{
   985  		{
   986  			name: "virtual-machines",
   987  			opts: &NearestOptions{
   988  				Type: "virtual",
   989  			},
   990  			want: true,
   991  		},
   992  		{
   993  			name: "sites",
   994  			opts: &NearestOptions{
   995  				Sites: []string{"foo"},
   996  			},
   997  			want: true,
   998  		},
   999  		{
  1000  			name: "none",
  1001  			opts: &NearestOptions{
  1002  				Type: "physical",
  1003  			},
  1004  			want: false,
  1005  		},
  1006  	}
  1007  
  1008  	for _, tt := range tests {
  1009  		t.Run(tt.name, func(t *testing.T) {
  1010  			got := alwaysPick(tt.opts)
  1011  			if got != tt.want {
  1012  				t.Errorf("alwaysPick() got: %v, want: %v", got, tt.want)
  1013  			}
  1014  		})
  1015  	}
  1016  }
  1017  
  1018  func TestPickWithProbability(t *testing.T) {
  1019  	tests := []struct {
  1020  		name        string
  1021  		probability float64
  1022  		seed        int64
  1023  		want        bool
  1024  	}{
  1025  		{
  1026  			name:        "no-probability",
  1027  			probability: 1.0,
  1028  			want:        true,
  1029  		},
  1030  		// If we use 2 as a seed, the pseudo-random number generated will be < 0.5.
  1031  		{
  1032  			name:        "pick-with-probability",
  1033  			probability: 0.5,
  1034  			seed:        2,
  1035  			want:        true,
  1036  		},
  1037  		// If we use 1 as a seed, the pseudo-random number generated will be > 0.5.
  1038  		{
  1039  			name:        "do-not-pick-with-probability",
  1040  			probability: 0.5,
  1041  			seed:        1,
  1042  			want:        false,
  1043  		},
  1044  	}
  1045  
  1046  	for _, tt := range tests {
  1047  		t.Run(tt.name, func(t *testing.T) {
  1048  			rand.Seed(tt.seed)
  1049  			got := pickWithProbability(tt.probability)
  1050  
  1051  			if got != tt.want {
  1052  				t.Errorf("pickWithProbability() got: %v, want: %v", got, tt.want)
  1053  			}
  1054  		})
  1055  	}
  1056  }
  1057  
  1058  func TestBiasedDistance(t *testing.T) {
  1059  	tests := []struct {
  1060  		name     string
  1061  		country  string
  1062  		r        *v2.Registration
  1063  		distance float64
  1064  		want     float64
  1065  	}{
  1066  		{
  1067  			name:    "empty-country",
  1068  			country: "",
  1069  			r: &v2.Registration{
  1070  				CountryCode: "foo",
  1071  			},
  1072  			distance: 100,
  1073  			want:     100,
  1074  		},
  1075  		{
  1076  			name:    "unknown-country",
  1077  			country: "ZZ",
  1078  			r: &v2.Registration{
  1079  				CountryCode: "foo",
  1080  			},
  1081  			distance: 100,
  1082  			want:     100,
  1083  		},
  1084  		{
  1085  			name:    "same-country",
  1086  			country: "foo",
  1087  			r: &v2.Registration{
  1088  				CountryCode: "foo",
  1089  			},
  1090  			distance: 100,
  1091  			want:     100,
  1092  		},
  1093  		{
  1094  			name:    "different-country",
  1095  			country: "bar",
  1096  			r: &v2.Registration{
  1097  				CountryCode: "foo",
  1098  			},
  1099  			distance: 100,
  1100  			want:     200,
  1101  		},
  1102  	}
  1103  
  1104  	for _, tt := range tests {
  1105  		t.Run(tt.name, func(t *testing.T) {
  1106  			got := biasedDistance(tt.country, tt.r, tt.distance)
  1107  
  1108  			if got != tt.want {
  1109  				t.Errorf("biasedDistance() got: %f, want: %f", got, tt.want)
  1110  			}
  1111  		})
  1112  	}
  1113  }