From d2396c88e16b9304afa41493df6b7445332cc9bb Mon Sep 17 00:00:00 2001
From: nkinkade <kinkade@measurementlab.net>
Date: Thu, 21 Nov 2024 09:39:35 -0700
Subject: [PATCH] Add new handler for the recently added siteinfo package
 (#209)

* 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.
---
 handler/handler.go                       | 28 ++++++++++++++
 handler/handler_test.go                  | 47 ++++++++++++++++++++++++
 heartbeat/heartbeattest/heartbeattest.go |  6 ++-
 locate.go                                |  3 ++
 openapi.yaml                             | 15 ++++++++
 5 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/handler/handler.go b/handler/handler.go
index 23606ff7..34cb6a00 100644
--- a/handler/handler.go
+++ b/handler/handler.go
@@ -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"
@@ -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) {
diff --git a/handler/handler_test.go b/handler/handler_test.go
index cc20fdb4..ddda464c 100644
--- a/handler/handler_test.go
+++ b/handler/handler_test.go
@@ -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 {
diff --git a/heartbeat/heartbeattest/heartbeattest.go b/heartbeat/heartbeattest/heartbeattest.go
index c5795983..cd65b8ad 100644
--- a/heartbeat/heartbeattest/heartbeattest.go
+++ b/heartbeat/heartbeattest/heartbeattest.go
@@ -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.
@@ -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
 }
 
diff --git a/locate.go b/locate.go
index 008c938e..5159f0cc 100644
--- a/locate.go
+++ b/locate.go
@@ -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,
diff --git a/openapi.yaml b/openapi.yaml
index 4c0e16c3..bb399b93 100644
--- a/openapi.yaml
+++ b/openapi.yaml
@@ -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: