Skip to content

Commit

Permalink
Add new handler for the recently added siteinfo package (#209)
Browse files Browse the repository at this point in the history
* Returns a list of all heartbeat registrations

* Adds /v2/siteinfo/registrations to openapi.yaml

* Adds new handler for the siteinfo package

* Adds 500 response to openapi spec for /v2/siteinfo/registrations

Also updates the description to be more accurate.
  • Loading branch information
nkinkade authored Nov 21, 2024
1 parent d0d5fac commit d2396c8
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 1 deletion.
28 changes: 28 additions & 0 deletions handler/handler.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"github.com/m-lab/locate/heartbeat"
"github.com/m-lab/locate/limits"
"github.com/m-lab/locate/metrics"
"github.com/m-lab/locate/siteinfo"
"github.com/m-lab/locate/static"
prom "github.com/prometheus/client_golang/api/prometheus/v1"
"github.com/prometheus/common/model"
Expand Down Expand Up @@ -229,6 +230,33 @@ func (c *Client) Ready(rw http.ResponseWriter, req *http.Request) {
}
}

// Registrations returns information about registered machines. There are 3
// supported query parameters:
//
// * format - defines the format of the returned JSON
// * org - limits results to only records for the given organization
// * exp - limits results to only records for the given experiment (e.g., ndt)
func (c *Client) Registrations(rw http.ResponseWriter, req *http.Request) {
var err error
var result interface{}

q := req.URL.Query()
format := q.Get("format")

switch format {
default:
result, err = siteinfo.Machines(c.LocatorV2.Instances(), q)
}

if err != nil {
v2Error := v2.NewError("siteinfo", err.Error(), http.StatusInternalServerError)
writeResult(rw, http.StatusInternalServerError, v2Error)
return
}

writeResult(rw, http.StatusOK, result)
}

// checkClientLocation looks up the client location and copies the location
// headers to the response writer.
func (c *Client) checkClientLocation(rw http.ResponseWriter, req *http.Request) (*clientgeo.Location, error) {
Expand Down
47 changes: 47 additions & 0 deletions handler/handler_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -319,6 +319,53 @@ func TestClient_Ready(t *testing.T) {
})
}
}
func TestClient_Registrations(t *testing.T) {
tests := []struct {
name string
instances map[string]v2.HeartbeatMessage
fakeErr error
wantStatus int
}{
{
name: "success-status-200",
instances: map[string]v2.HeartbeatMessage{
"ndt-mlab1-abc0t.mlab-sandbox.measurement-lab.org": {},
},
wantStatus: http.StatusOK,
},
{
name: "error-status-500",
instances: map[string]v2.HeartbeatMessage{
"invalid-hostname.xyz": {},
},
fakeErr: errors.New("fake error"),
wantStatus: http.StatusInternalServerError,
},
}
for _, tt := range tests {
fakeStatusTracker := &heartbeattest.FakeStatusTracker{
Err: tt.fakeErr,
FakeInstances: tt.instances,
}

t.Run(tt.name, func(t *testing.T) {
c := NewClient("foo", &fakeSigner{}, &fakeLocatorV2{StatusTracker: fakeStatusTracker}, nil, nil, nil)

mux := http.NewServeMux()
mux.HandleFunc("/v2/siteinfo/registrations/", c.Registrations)
srv := httptest.NewServer(mux)
defer srv.Close()

req, err := http.NewRequest(http.MethodGet, srv.URL+"/v2/siteinfo/registrations?org=mlab", nil)
rtx.Must(err, "failed to create request")
resp, err := http.DefaultClient.Do(req)
rtx.Must(err, "failed to issue request")
if resp.StatusCode != tt.wantStatus {
t.Errorf("Registrations() wrong status; got %d; want %d", resp.StatusCode, tt.wantStatus)
}
})
}
}

func TestExtraParams(t *testing.T) {
tests := []struct {
Expand Down
6 changes: 5 additions & 1 deletion heartbeat/heartbeattest/heartbeattest.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,8 @@ func (c *fakeErrorMemorystoreClient[V]) GetAll() (map[string]V, error) {

// FakeStatusTracker provides a fake implementation of HeartbeatStatusTracker.
type FakeStatusTracker struct {
Err error
Err error
FakeInstances map[string]v2.HeartbeatMessage
}

// RegisterInstance returns the FakeStatusTracker's Err field.
Expand All @@ -73,6 +74,9 @@ func (t *FakeStatusTracker) UpdatePrometheus(hostnames, machines map[string]bool

// Instances returns nil.
func (t *FakeStatusTracker) Instances() map[string]v2.HeartbeatMessage {
if t.FakeInstances != nil {
return t.FakeInstances
}
return nil
}

Expand Down
3 changes: 3 additions & 0 deletions locate.go
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,9 @@ func main() {
mux.HandleFunc("/v2/live", c.Live)
mux.HandleFunc("/v2/ready", c.Ready)

// Return list of all heartbeat registrations
mux.HandleFunc("/v2/siteinfo/registrations", c.Registrations)

srv := &http.Server{
Addr: ":" + listenPort,
Handler: mux,
Expand Down
15 changes: 15 additions & 0 deletions openapi.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,21 @@ paths:
tags:
- platform

"/v2/siteinfo/registrations":
get:
description: |-
Returns heartbeat registration information in various formats.
operationId: "v2-siteinfo-registrations"
produces:
- "application/json"
responses:
'200':
description: OK.
'500':
description: Error.
tags:
- siteinfo

definitions:
# Define the query reply without being specific about the structure.
ErrorResult:
Expand Down

0 comments on commit d2396c8

Please sign in to comment.